jax

4a5e0e84 - Raise an error if we see an invalid memory kind in device_put abstract_eval