Add support for Eigen's Tensor module

[ManifoldFR]

This PR adds support for the types in Eigen's "unsupported" Tensor module, from the header <unsupported/Eigen/CXX11/Tensor> (we keep CXX11/ in the include to support Eigen3).

To-do list

• Plain Eigen::Tensor types
• Eigen::TensorMap<T>
• Eigen::TensorRef<T>
• Tests
  • Add/multiply plain Tensor types
  • Plain Tensor as struct member

Background

We are contributing this upstream as it is part of the porting of Pinocchio over to nanobind 😄

[wjakob]

Hi @ManifoldFR,

this looks great!

I responded to one comment inline. A few other thoughts:

Test coverage seems a bit less than the tensor module of pybind11. I recommend taking a look at that and recovering useful bits.
https://github.com/pybind/pybind11/blob/master/tests/test_eigen_tensor.py

Are you testing special cases like tensors with length zero along an axis, and scalars (ndim=0)?

Why are there GC calls in the test suite? I don't think you are checking that things are actually deallocated.

Thanks,
Wenzel

[ManifoldFR]

Thanks for the feedback Wenzel.

I can't see your reply in the comment? Might be a GitHub bug.

I didn't realize pybind11 had official Tensor support, I'll checked out their code.

As for the gc test, I merely replicated what's done in the current Eigen dense test suite:

nanobind/tests/test_eigen.py

Lines 207 to 223 in cfbb02f

	def test06_map():
	b = t.Buffer()
	m = b.map()
	dm = b.dmap()
	for i in range(10):
	for j in range(3):
	m[i, j] = i*3+j
	for i in range(10):
	for j in range(3):
	assert dm[i, j] == i*3+j
	del dm
	del b
	gc.collect()
	gc.collect()
	for i in range(10):
	for j in range(3):
	assert m[i, j] == i*3+j

I wrote a few more comprehensive tests (casts, 0-dim tensor and tensors with zeroes in one of the sizes).

[wjakob]

Forgot to submit the review.

[wjakob]

Thank you @ManifoldFR. One more thing: can you explanation on how to use the Eigen type casters to docs/eigen.rst?

[wjakob]

I think you can remove the gc.collect() calls from the Eigen test suite. They are likely leftovers from a debugging session where I checked for leaks at the end.

[wjakob]

Was this meant to be a .. remark RST block? Probably it's overkill, you can just merge that with the previous paragraph. Are non-contiguous arrays just unsupported for Map/Ref or also for the base tensor?

[ManifoldFR]

Non-contiguous arrays are not supported for tensor maps/refs nor the base tensor type, because these types themselves are contiguous in Eigen. I'll edit the text.

[wjakob]

A question about that: if you pass a non-contiguous matrix or vector to a function taking Eigen::Matrix or Eigen::Vector, nanobind will AFAIK make a contiguous copy if implicit conversions are allowed. This would be a surprising deviation from the vector/matrix case. nanobind's underlying ndarray infrastructure can do this automatically, so I'm wondering where the difference comes from (and if it is intentional)

[ManifoldFR]

You're right, I just tested taking a slice on a 3d tensor and passing it to a function taking a const reference to a plain tensor type, and it worked; I'll add a test case.

[wjakob]

Sounds good. In that case, this comment is moot, right? "Note that taking
non-contiguous NumPy arrays as arguments is not supported for any of these types."

It would be also worth adding a test (if not already present) that Map and Ref (less sure about this one) refuse such a tensor. At least for Map, the idea is IIRC that it wraps a data buffer without copies, so implicit conversions are very much not wanted.

[ManifoldFR]

I've added the relevant pytest.raises tests and updated the comment in eigen.rst.

I know that TensorRef can hold a reference to the striding expression/operation but that needs a base Eigen Tensor/TensorMap to start from. Perhaps someone could do something with this (with some kind of custom op replicating it with an external pointer?)

[wjakob]

I left one more comment. All good after this? (i.e., ready to merge?)

[ManifoldFR]

I left one more comment. All good after this? (i.e., ready to merge?)

After this, all good :)

[wjakob]

Thank you for your hard work @ManifoldFR!

[ManifoldFR]

Thanks for merging this!! This will be useful for our big upgrade to nanobind 😄