Adam Paszke 9c5e3f7ecc Verify that slices are trivial before discarding them in state primitives ...

At the moment, if `r` is a JAX ref then `r[0:1] = a` works, but it silently ignores the slices
and performs `r[:] = a` instead...

PiperOrigin-RevId: 529385973

2023-05-04 05:59:47 -07:00

..

__init__.py

Add initial implementation of a run_state primitive

2023-04-03 21:32:32 -07:00

discharge.py

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

primitives.py

Verify that slices are trivial before discarding them in state primitives

2023-05-04 05:59:47 -07:00

types.py

Merge pull request #15149 from sharadmv:runstate

2023-04-04 10:56:25 -07:00

utils.py

Add initial implementation of a run_state primitive

2023-04-03 21:32:32 -07:00