jax

0c543aef - Match the argument name with the name in `Args` section in docstring