Matthew Johnson e0d2736e37 add custom_jvp for jax.nn.softmax ...

This avoids saving the jnp.exp(...) value.

2023-04-22 11:28:03 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

functions.py

add custom_jvp for jax.nn.softmax

2023-04-22 11:28:03 -07:00

initializers.py

Add test framework for module attribute

2023-04-21 13:20:16 -07:00