Type Caster for linear algebra library Armadillo (CPU) and Bandicoot (GPU)

[r1cheu]

Armadillo and Bandicoot offer nearly identical syntax for linear algebra operations, making it easy to transition between CPU and GPU implementations.

It would be great if nanobind could support conversions between Python libraries like NumPy, JAX, etc., and these libraries directly.

[wjakob]

nanobind offers the ndarray interface that allows it to accept and return general nd-arrays from CPU/GPU frameworks. This could be used to implement bindings for libraries like armadillo or bandicoot. I suggest that you review the documentation of the ndarray feature.

[theAeon]

thanks for quick reply. I’ll dig into the ndarray interface. I’d willing to implement a type caster, like the one for Eigen(though I’m quite new to template metaprogramming, and it may take quite a long time). Any advice?

Hi-I have an implementation modeled on the sparse/dense eigen casters here if its of any use.

[r1cheu]

Hi-I have an implementation modeled on the sparse/dense eigen casters here if its of any use.

Helps a lot. Previously, I directly implemented to_py and to_arma to cast matrix manually, maybe now is the time to make the bind code cleaner 😀.

[theAeon]

Would definitely appreciate any patches if you happen to tweak it!

[r1cheu]

hmmm, armadillo provides constructor mat(ptr_aux_mem, n_rows, n_cols, copy_aux_mem = false, strict = true), can be combined with keep_alive to avoid copy when converting from numpy to arma, if we can make sure the mat shape is not changed, which should be a common case (at least in my case), but I'm not very clear about how to cast for arma::Col, arma::Mat, and arma::Cube in one type caster, or if I should implement three of them. I'll see what I can do tomorrow.

[r1cheu]

hmmm, armadillo provides constructor mat(ptr_aux_mem, n_rows, n_cols, copy_aux_mem = false, strict = true), can be combined with keep_alive to avoid copy when converting from numpy to arma, if we can make sure the mat shape is not changed, which should be a common case (at least in my case), but I'm not very clear about how to cast for arma::Col, arma::Mat, and arma::Cube in one type caster, or if I should implement three of them. I'll see what I can do tomorrow.

  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {

    make_caster<NDArray> caster;
    if (!caster.from_python(src, flags & ~(uint8_t)cast_flags::accepts_none,
                            cleanup))
      return false;

    NDArray &t = caster.value;
    if constexpr (ndim_v<T> == 1) {
      value = T(t.data(), t.shape(0), false, true);
    } else {
      value = T(t.data(), t.shape(0), t.shape(1), false, true);
    }
    return true;
  }

This code snippet can be used to achieve zero-copy. If needed, use it with nb::keep_alive to ensure the memory won't be deallocated.