Closed copybara-service[bot] closed 6 months ago
Avoid multiplication of boolean arrays
Instead, use a jnp.where statement. That avoids type promotion of boolean arrays to integer arrays.
Fixes issue https://github.com/google-deepmind/optax/issues/828
Also made cosmetic changes to the docstring
Avoid multiplication of boolean arrays
Instead, use a jnp.where statement. That avoids type promotion of boolean arrays to integer arrays.
Fixes issue https://github.com/google-deepmind/optax/issues/828
Also made cosmetic changes to the docstring