google / deluca

Performant, differentiable reinforcement learning
https://deluca.fyi
Apache License 2.0
121 stars 17 forks source link

[JAX] Replace uses of `jnp.array` in types with `jnp.ndarray`. #72

Open copybara-service[bot] opened 1 year ago

copybara-service[bot] commented 1 year ago

[JAX] Replace uses of jnp.array in types with 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.