google / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Home Page:http://jax.readthedocs.io/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

`jnp.take` out of bound documentation incorrect

HanGuo97 opened this issue · comments

Please:

  • Check for duplicate issues.
  • Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:
jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0
)
# DeviceArray([0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4], dtype=int32)

jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0,
    mode="raise",
)
# NotImplementedError

According to the documentation, the default out-of-bound behavior is "raise". However, this is not the case in practice

Thanks for the report - I'll work on a documentation fix.