feat: add support for `jax` like `to_torch`

[austinv11]

Description of feature

Excellent work, I was just wondering if there are any plans to support JAX tensors in the future? I have recently started using JAX more often than PyTorch for better performance.

Thanks!

[ilan-gold]

@austinv11 You've actually touched on one of my favorite things about this library - we don't depend on torch so you can just "use jax".

Concretely, when you yield your data out of Loader.__iter__ and have to_torch=False, you get back one of the types listed (scipy/cupy sprase matrix or numpy/cupy dense array), which you can then wrap in the appropriate jax data structure:

1. If preload_to_gpu=True and you get back a sparse cupy matrix, you can densify the sparse matrix output and put the result 0-copy into a jax array via dlpack
2. If preload_to_gpu=False and you get back a scipy.sparse.matrix, you can create a jax sparse datastructure and then densify on the GPU (probably faster than the CPU, but haven't tried)
3. If your result is dense, you can just use jax.array(...) or similar.

I'm going to rename the issue to add support for this logic internally. We can make a breaking change, perhaps, to migrate from to_torch to to: Literal["jax", "torch"]. Would love a contribution here if you're up for it, but I also don't think this is too complicated. LMK!

Thanks for the issue!

[ilan-gold]

P.S We have GPU CI so this is very testable :)

[austinv11]

Awesome! I'll see if I can whip up a PR

[ilan-gold]

Thanks - I should have been clear also about how the breaking change would work. I would leave the to_torch argument for now and deprecate its use:

1. Give it a new default value like the string auto so to_torch: Literal["auto"] | bool = "auto"
2. encode the current behavior of to_torch=True if torch is installed behind a check for if to_torch == "auto"
3. Warn if a user has set to_torch to a boolean explicitly to stop doing that and use to: Literal["jax", "torch"] | None instead

Basically, we point people to a new to: Literal["jax", "torch"] | None = None argument if they try to set to_torch.

I think we can then use jax or torch by default if they are installed. I don't think we really have to worry about environments with both from what I understand. It is apparently basically impossible to install both together, right?