jax

0ff0d7b9 - jnp.take: fix annotation for fill_value