JAX v0.5.3

• New Features

  • Added a 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. If set to false, JAX does not need to
    emit code clamping negative indices, which improves code size.
  • Added a replace option to jax.random.categorical to enable sampling
    without replacement.