Makes me wonder if we shouldn't do this in more places? Could also be worth checking if it works for Jax and TF in that case. If we make it work for all of those and can make it work more broadly that would be very worth documenting and bragging about!

By do this you mean leverage dlpack?

I can manually check if this works for JAX and TF (it would be weird if it didn't) but adding the unit test requires a bit more infra since we don't have a separate test env and those are heavy packages to import

By do this you mean leverage dlpack?

Yes, I mean have the ability to convert anything with a __dlpack__ attribute to an arrow array. I think I may have added some numpy conversion logic that was supposed to work with pytorch earler but it may have been removed since.

If I don't miss-remember there should also be some way to appease mypy and friends by annotating a type hint of the style "anything that has this attribute" but not sure.

I can manually check if this works for JAX and TF (it would be weird if it didn't) but adding the unit test requires a bit more infra since we don't have a separate test env and those are heavy packages to import

Yeah manual check seems fine.

Confirmed it works for jax:

    elif version == "raw_jax":
        timing = simple_loop(shape, raw_any, partial(jax.random.uniform, jax.random.key(123)))

>>>python sample.py raw_jax
Loop 100
1
raw_jax:
        0.0025s

Yes, I mean have the ability to convert anything with a dlpack attribute to an arrow array. I think I may have added some numpy conversion logic that was supposed to work with pytorch earler but it may have been removed since.

IIUC we can convert any 1d object to a pyarrow array with this. Anything greater than 1d either needs to be flattened which could still be fast or made a nested list of lists as a pyarrow limitation

If I don't miss-remember there should also be some way to appease mypy and friends by annotating a type hint of the style "anything that has this attribute" but not sure.

Yes there is. We should be able to use protocols for that

Tensorflow doesn't work yet because they haven't fully committed to the standard yet
https://www.tensorflow.org/api_docs/python/tf/experimental/dlpack/to_dlpack

But on main it appears they might support __dlpack__ which landed 2 days after the next last release. So on the next release it hopefully works. tensorflow/tensorflow#89079