Change8

jax-v0.5.3

📦 jaxView on GitHub →
2 features🔧 3 symbols

Summary

This release introduces new options for slicing functions and categorical sampling in JAX, improving code size and adding support for sampling without replacement.

✨ New Features

  • Added an `allow_negative_indices` option to `jax.lax.dynamic_slice`, `jax.lax.dynamic_update_slice` and related functions. The default is true, matching the current behavior. Setting it to false removes code that clamps negative indices, reducing code size.
  • Added a `replace` option to `jax.random.categorical` to enable sampling without replacement.

🔧 Affected Symbols

jax.lax.dynamic_slicejax.lax.dynamic_update_slicejax.random.categorical