Matthew Johnson a964dc3b9a simpler pretty-print for pjit, tweak custom pp rule signature

2023-02-09 12:45:51 -08:00

..

__init__.py

Enable state effect in cond_p (except in grad and vmap)

2022-11-02 16:07:58 -07:00

discharge.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

primitives.py

simpler pretty-print for pjit, tweak custom pp rule signature

2023-02-09 12:45:51 -08:00

types.py

Enable state effect in cond_p (except in grad and vmap)

2022-11-02 16:07:58 -07:00