Closed copybara-service[bot] closed 1 year ago
[JAX] Replace uses of jnp.array in types with jnp.ndarray.
jnp.array
jnp.ndarray
jnp.array is a function, not a type: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.array.html so it never makes sense to use jnp.array in a type annotation.
Presumably the intent was to write jnp.ndarray aka jax.Array. Change uses of jnp.array to jnp.ndarray.
jax.Array
[JAX] Replace uses of
jnp.arrayin types withjnp.ndarray.jnp.arrayis a function, not a type: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.array.html so it never makes sense to usejnp.arrayin a type annotation.Presumably the intent was to write
jnp.ndarrayakajax.Array. Change uses ofjnp.arraytojnp.ndarray.