jax-ml/jax jax-v0.5.3

JAX v0.5.3

on GitHub

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.