From 4c7b3096e5362d87458b8ff2ef0b0b2720ae8c8a Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Tue, 17 Oct 2023 13:03:32 -0700 Subject: [PATCH] [LSC] change uses of jax.random.KeyArray and jax.random.PRNGKeyArray to jax.Array This change replaces uses of jax.random.KeyArray and jax.random.PRNGKeyArray in the context of type annotations with jax.Array, which is the correct annotation for JAX PRNG keys moving forward. The purpose of this change is to remove references to KeyArray and PRNGKeyArray, which are deprecated (https://github.com/google/jax/pull/17594) and will soon be removed from JAX. The design and thought process behind this is described in https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html. Note that KeyArray and PRNGKeyArray have always been aliased to Any, so the new type annotation is far more specific than the old one. PiperOrigin-RevId: 574241296 --- learned_optimization/optimizers/optax_opts.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/learned_optimization/optimizers/optax_opts.py b/learned_optimization/optimizers/optax_opts.py index e8c1edd..8df1a53 100644 --- a/learned_optimization/optimizers/optax_opts.py +++ b/learned_optimization/optimizers/optax_opts.py @@ -384,11 +384,13 @@ def __init__(self, # SM3 doesn't support scalars, so we have to reshape the params and grads. - def init(self, - params: Any, - model_state: Optional[Any] = None, - num_steps: Optional[int] = None, - key: chex.PRNGKey = None) -> SM3OptState: + def init( + self, + params: Any, + model_state: Optional[Any] = None, + num_steps: Optional[int] = None, + key: Optional[chex.PRNGKey] = None, + ) -> SM3OptState: should_reshape = jax.tree_util.tree_map(lambda x: len(x.shape) == 0, params) # pylint: disable=g-explicit-length-test params = jax.tree_util.tree_map(_expand_scalar, params, should_reshape) out = super().init(params, model_state, num_steps, key)