Stride-related surprises with the Eigen casters

[oremanj]

I've been looking into nanobind's <nanobind/eigen/dense.h> type casters recently, and noticed a couple of surprising behaviors. They aren't quite bugs, but they seem like they might be valuable to improve upon.

One: The caster for Eigen plain types (Matrix, etc) requires that the input array be able to be converted into the memory layout that Eigen expects. However, this is generally Fortran-contiguous, and the non-numpy array types that nanobind supports (jax, tensorflow, pytorch) AFAICT provide no way to produce a Fortran-contiguous array. (You might be able to get pytorch to do it with an appropriate transposition, but eg tensorflow doesn't have a concept of variable strides at all.) Thus, if you pass in one of those other frameworks' array objects, the cast to the Eigen type will fail. This seems like an unnecessary limitation given that we're making a copy anyway.

Two: When casting to Eigen::Ref<const T> and a dtype conversion is required, in most cases two copies will be made: one to make an ndarray with the correct dtype (but any strides), and another in the constructor of Eigen::Ref to get the data to be laid out with the correct strides. Similarly to the previous point, it feels like if we're making a copy we should be able to make it with the correct layout.

Both of these seem like consequences of the approach "cast to an appropriate ndarray, then wrap that", given that the conversion is always done using methods provided by the exporter of the original array data. I'm not sure the best way to resolve them. For plain types, nanobind could optionally support inputs with different shape than the output, copying element-by-element rather than doing a memcpy; for Refs, nanobind could use the non-DStrided Map type if the input can be put in that form; but I'm not sure if either of these solutions is ideal, or worth the complexity cost.

[wjakob]

I don't use the Eigen binding code myself. Let me add @WKarel who has contributed major portions of it (in particular, the bits that I think you are referring to).

[WKarel]

1. the failure in your use case could be avoided in ndarray_import: if the required device is the CPU (true when converting to Eigen), and if NumPy is available (may not be the case), then np.asarray could be used to create a view or copy of the source array, with the required data type and memory layout. With torch.tensor, this copies the underlying data if the required layout is column-major, but not if it is row-major. I do not know to which extent this method would work for other array packages.
   However, the above would still make 2 copies instead of 1 in more situations than necessary. To minimize that number, type_caster<PlainEigen>::from_python would need to cast to a nb::ndarray<any_contig> (but of the required shape and data type, possibly allowing type conversions). And then consider the ndarray's strides when copying. That should be doable quite easily, possibly reusing the type_caster for Eigen::Map.
2. type_caster<Eigen::Ref<T>, ...>::from_python first tries to map the array with the required type and strides, without copying the data. If that fails and T is const (and Ref can map contiguous memory), then it tries to map the array or a copy thereof with the required type and fully dynamic strides. Indeed, if the Ref to be returned does not support mapping these strides, then this results in another copy being made. This is the case if an array is passed with mismatching type and memory layout. The second copy could be avoided if the second caster tried to cast to either f_contig or c_contig instead of any_contig, if the target Ref is able to reference that memory layout without copying data. It might actually be enough to only change the definition of type_caster::DMap. And I would leave any_contig for the general case, so copies are avoided if Ref requires unusual strides that the passed array happens to match.

By the way, ndarray_import makes a copy when converting a np.ndarray for which only the C- or F-order does not match, which could be avoided by using np.asarray instead of np.ndarray.astype (which always copies).