Use dlpack for tensor export again, as referenced issue is resolved

[adconner]

Upstream issue resolved, so we can use the non-copying export.

[samuela]

sorry, ci should be fixed after b958825. would you mind rebasing?

[samuela]

hmm looks like this requires detaching before conversion

[adconner]

Okay I was mistaken about the problem being fixed, was mainly going by the linked jax issue jax-ml/jax#25066 being resolved (did the test for nix build, which isn't the full test suite and didn't catch the failuer). Test suite is still failing even after detaching and contiguous and trying from_dlpack(..,copy=True). Hard to understand what can be causing the error. Perhaps a new issue needs to be filed with upstream if they think this is resolved.

I could file an issue here if you wish, but without using dlpack, in addition to causing out of memory errors on my problem, it doesn't handle bfloat16 transfer since numpy doesn't and never will. So I think we are in a state where something will be wrong no matter what (and on my side I will certainly be using the dlpack interface anyway since I dont run into this Unexpected XLA layout problem).

[samuela]

Ahh, I see. That's unfortunate. Well thank you for exploring and describing the problem in detail. I had not previously considered the bf16 case... that's a good point

[adconner]

A couple of things, it seems like the version of jax used in the flake is old enough to not get the fix from XLA (0.5.0), so maybe this issue actually fixed in a new version. In any case, I looked at the upsteam issue again and realized it is easy to work around even in old jax. Here's an updated, non copying implementation passing the tests and which works on bfloat16. Once you want to migrate to recent jax, you might be able to remove this workaround if you want, but its harmless.

[samuela]

oh cool! do you happen to know the minimum jax version necessary to avoid the workaround?

[samuela]

ok i just pushed 102c45f which updates the jax version to 0.6.2

[adconner]

No I did not track this down. Its possible its in the future. Presumably you can binary search the answer changing the version in your flake and running the tests.

[samuela]

okok would you mind rebasing onto the latest main and seeing if we can get rid of the workaround? then we can merge

[adconner]

looks good

[samuela]

Thanks @adconner !