Commits

Commits on Aug 23, 2023

Replace uses of `jnp.array` in types with `jnp.ndarray`.

hawkinsp
authored and
OptaxDev
committed

Commits on Aug 14, 2023

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

hawkinsp
authored and
OptaxDev
committed
[JAX] Fix incorrect type annotations.

hawkinsp
authored and
OptaxDev
committed

Commits on Apr 20, 2023

Fix pytype failures related to teaching pytype about NumPy scalar types.

hawkinsp
authored and
OptaxDev
committed

Commits on Mar 14, 2023

[JAX] Fix some type annotations in optax that break with an upcoming JAX change.

hawkinsp
authored and
OptaxDev
committed

Commits on Mar 11, 2023

Replaces references to jax.numpy.DeviceArray with jax.Array.

hawkinsp
authored and
OptaxDev
committed

Commits on Feb 18, 2023

[JAX] Fix pytype error in Optax if jnp.ndarray == jax.Array.

hawkinsp
authored and
OptaxDev
committed

Commits on Feb 14, 2023

[JAX] Replace uses of jax.xla_computation() with jax.jit().lower().

hawkinsp
authored and
OptaxDev
committed

Commits on Oct 7, 2022

[JAX] Import scipy submodules directly in code that uses them.

hawkinsp
authored and
OptaxDev
committed

Commits on Oct 4, 2021

[JAX] Increase numerical tolerances of tests in preparation for an XLA:CPU vectorization change.

hawkinsp
authored and
OptaxDev
committed