Exciting! I have been wanting a Rust port of JAX for a long time! (I mostly care about the ability to write numerical GPU code via XLA, not differentiation)

Is there a crate in the works? I would start using it (and contributing if needed) instantly.