jax

fa383b4a - Mark primitive parameters as keyword-only arguments in rules in lax.py. (#2625)

Mark primitive parameters as keyword-only arguments in rules in lax.py. (#2625) * Mark primitive parameters as keyword-only arguments in rules in lax.py. * Fix dynamic update slice batching rule. * Fix dynamic slice batching rule.