jax.experimental.custom_dce
module#
API#
|
Customize the DCE behavior of a JAX-transformable function. |
|
Define a custom DCE rule for this function. |
jax.experimental.custom_dce
module#
|
Customize the DCE behavior of a JAX-transformable function. |
|
Define a custom DCE rule for this function. |