revert rollback of seed key v3 PR #1240

PiperOrigin-RevId: 748434720

Author: vroulet

optax/_src/alias.py

@@ -17,14 +17,18 @@
 from collections.abc import Callable
 import functools
 from typing import Any, Optional, Union
+import warnings
 
+import chex
+import jax
 import jax.numpy as jnp
 from optax._src import base
 from optax._src import clipping
 from optax._src import combine
 from optax._src import factorized
 from optax._src import linesearch as _linesearch
 from optax._src import transform
+from optax._src import utils
 from optax._src import wrappers
 
 
@@ -1280,7 +1284,7 @@ def noisy_sgd(
     learning_rate: base.ScalarOrSchedule,
     eta: float = 0.01,
     gamma: float = 0.55,
-    seed: int = 0,
+    key: chex.PRNGKey | int | None = None
 ) -> base.GradientTransformationExtraArgs:
   r"""A variant of SGD with added noise.
 
@@ -1311,7 +1315,7 @@ def noisy_sgd(
     eta: Initial variance for the Gaussian noise added to gradients.
     gamma: A parameter controlling the annealing of noise over time ``t``, the
       variance decays according to ``(1+t)**(-gamma)``.
-    seed: Seed for the pseudo-random generation process.
+    key: a PRNG key used as the random key.
 
   Returns:
     The corresponding :class:`optax.GradientTransformationExtraArgs`.
@@ -1321,7 +1325,10 @@ def noisy_sgd(
     >>> import jax
     >>> import jax.numpy as jnp
     >>> def f(x): return jnp.sum(x ** 2)  # simple quadratic function
-    >>> solver = optax.noisy_sgd(learning_rate=0.003)
+    >>> solver = optax.noisy_sgd(
+    ...   learning_rate=0.003,
+    ...   key=jax.random.key(0)
+    ... )
     >>> params = jnp.array([1., 2., 3.])
     >>> print('Objective function: ', f(params))
     Objective function:  14.0
@@ -1341,8 +1348,13 @@ def noisy_sgd(
     Neelakantan et al, `Adding Gradient Noise Improves Learning for Very Deep
     Networks <https://arxiv.org/abs/1511.06807>`_, 2015
   """
+  if key is None:
+    warnings.warn(
+        'Specifying a key for optax.noisy_sgd will be required in optax 0.3.0.'
+    )
+    key = jax.random.key(0)
   return combine.chain(
-      transform.add_noise(eta, gamma, seed),
+      transform.add_noise(eta, gamma, key=utils.to_random_key(key)),
       transform.scale_by_learning_rate(learning_rate),
   )
 

optax/_src/alias_test.py

@@ -71,7 +71,7 @@
     {'opt_name': 'nadamw', 'opt_kwargs': {'learning_rate': 1e-2}},
     {
         'opt_name': 'noisy_sgd',
-        'opt_kwargs': {'learning_rate': 1e-3, 'eta': 1e-4},
+        'opt_kwargs': {'learning_rate': 1e-3, 'key': 0, 'eta': 1e-4},
     },
     {'opt_name': 'novograd', 'opt_kwargs': {'learning_rate': 1e-3}},
     {'opt_name': 'optimistic_adam', 'opt_kwargs': {'learning_rate': 2e-3}},

optax/_src/float64_test.py

@@ -48,14 +48,14 @@
         {'step_size_fn': lambda x: x * 0.1},
     ),
     ('scale_by_trust_ratio', transform.scale_by_trust_ratio, {}),
-    ('add_noise', transform.add_noise, {'eta': 1.0, 'gamma': 0.1, 'seed': 42}),
+    ('add_noise', transform.add_noise, {'key': 42, 'eta': 1.0, 'gamma': 0.1}),
     ('apply_every_k', transform.apply_every, {}),
     ('adagrad', alias.adagrad, {'learning_rate': 0.1}),
     ('adam', alias.adam, {'learning_rate': 0.1}),
     ('adamw', alias.adamw, {'learning_rate': 0.1}),
     ('fromage', alias.fromage, {'learning_rate': 0.1}),
     ('lamb', alias.lamb, {'learning_rate': 0.1}),
-    ('noisy_sgd', alias.noisy_sgd, {'learning_rate': 0.1}),
+    ('noisy_sgd', alias.noisy_sgd, {'learning_rate': 0.1, 'key': 0}),
     ('rmsprop', alias.rmsprop, {'learning_rate': 0.1}),
     ('sgd', alias.sgd, {'learning_rate': 0.1}),
     ('sign_sgd', alias.sgd, {'learning_rate': 0.1}),

optax/_src/utils.py

@@ -46,6 +46,15 @@ def canonicalize_dtype(
   return dtype
 
 
+def to_random_key(key_or_seed: chex.PRNGKey | int) -> chex.PRNGKey:
+  """Canonicalize a random key or an int representing a seed to a random key."""
+  if (isinstance(key_or_seed, jax.Array) and jnp.issubdtype(
+      key_or_seed.dtype, jax.dtypes.prng_key
+  )):
+    return key_or_seed
+  return jax.random.key(key_or_seed)
+
+
 @functools.partial(
     chex.warn_deprecated_function, replacement='optax.tree_utils.tree_cast'
 )
@@ -106,10 +115,10 @@ def __init__(self, loc: chex.Array, log_scale: chex.Array):
         self._mean.shape, self._scale.shape
     )
 
-  def sample(self, shape: Sequence[int], seed: chex.PRNGKey) -> chex.Array:
+  def sample(self, shape: Sequence[int], key: chex.PRNGKey) -> chex.Array:
     sample_shape = tuple(shape) + self._param_shape
     return (
-        jax.random.normal(seed, shape=sample_shape) * self._scale + self._mean
+        jax.random.normal(key, shape=sample_shape) * self._scale + self._mean
     )
 
   def log_prob(self, x: chex.Array) -> chex.Array:

optax/contrib/_common_test.py

@@ -105,7 +105,7 @@
     {'opt_name': 'lion', 'opt_kwargs': {'learning_rate': 1.0, 'b1': 0.99}},
     {
         'opt_name': 'noisy_sgd',
-        'opt_kwargs': {'learning_rate': 1.0, 'eta': 1e-4},
+        'opt_kwargs': {'learning_rate': 1.0, 'key': 0, 'eta': 1e-4},
     },
     {'opt_name': 'novograd', 'opt_kwargs': {'learning_rate': 1.0}},
     {

optax/contrib/_privacy.py

@@ -16,11 +16,13 @@
 
 from typing import Any, NamedTuple, Optional
 
+import chex
 import jax
 from optax._src import base
 from optax._src import clipping
 from optax._src import combine
 from optax._src import transform
+from optax._src import utils
 
 
 class DifferentiallyPrivateAggregateState(NamedTuple):
@@ -33,14 +35,16 @@ class DifferentiallyPrivateAggregateState(NamedTuple):
 
 
 def differentially_private_aggregate(
-    l2_norm_clip: float, noise_multiplier: float, seed: int
+    l2_norm_clip: float,
+    noise_multiplier: float,
+    key: chex.PRNGKey | int
 ) -> base.GradientTransformation:
   """Aggregates gradients based on the DPSGD algorithm.
 
   Args:
     l2_norm_clip: maximum L2 norm of the per-example gradients.
     noise_multiplier: ratio of standard deviation to the clipping norm.
-    seed: initial seed used for the jax.random.PRNGKey
+    key: a PRNG key used as the random key.
 
   Returns:
     A :class:`optax.GradientTransformation`.
@@ -63,10 +67,11 @@ def differentially_private_aggregate(
     specific way.
   """
   noise_std = l2_norm_clip * noise_multiplier
+  key = utils.to_random_key(key)
 
   def init_fn(params):
     del params
-    return DifferentiallyPrivateAggregateState(rng_key=jax.random.PRNGKey(seed))
+    return DifferentiallyPrivateAggregateState(rng_key=utils.to_random_key(key))
 
   def update_fn(updates, state, params=None):
     del params
@@ -91,7 +96,7 @@ def dpsgd(
     learning_rate: base.ScalarOrSchedule,
     l2_norm_clip: float,
     noise_multiplier: float,
-    seed: int,
+    key: chex.PRNGKey | int,
     momentum: Optional[float] = None,
     nesterov: bool = False,
 ) -> base.GradientTransformation:
@@ -106,7 +111,7 @@ def dpsgd(
     learning_rate: A fixed global scaling factor.
     l2_norm_clip: Maximum L2 norm of the per-example gradients.
     noise_multiplier: Ratio of standard deviation to the clipping norm.
-    seed: Initial seed used for the jax.random.PRNGKey
+    key: a PRNG key used as the random key.
     momentum: Decay rate used by the momentum term, when it is set to `None`,
       then momentum is not used at all.
     nesterov: Whether Nesterov momentum is used.
@@ -133,7 +138,7 @@ def dpsgd(
       differentially_private_aggregate(
           l2_norm_clip=l2_norm_clip,
           noise_multiplier=noise_multiplier,
-          seed=seed,
+          key=utils.to_random_key(key),
       ),
       (
           transform.trace(decay=momentum, nesterov=nesterov)

optax/contrib/_privacy_test.py

@@ -45,7 +45,9 @@ def setUp(self):
   def test_no_privacy(self):
     """l2_norm_clip=MAX_FLOAT32 and noise_multiplier=0 should recover SGD."""
     dp_agg = _privacy.differentially_private_aggregate(
-        l2_norm_clip=jnp.finfo(jnp.float32).max, noise_multiplier=0.0, seed=0
+        l2_norm_clip=jnp.finfo(jnp.float32).max,
+        noise_multiplier=0.0,
+        key=jax.random.key(0)
     )
     state = dp_agg.init(self.params)
     update_fn = self.variant(dp_agg.update)
@@ -59,7 +61,9 @@ def test_no_privacy(self):
   @parameterized.parameters(0.5, 10.0, 20.0, 40.0, 80.0)
   def test_clipping_norm(self, l2_norm_clip):
     dp_agg = _privacy.differentially_private_aggregate(
-        l2_norm_clip=l2_norm_clip, noise_multiplier=0.0, seed=42
+        l2_norm_clip=l2_norm_clip,
+        noise_multiplier=0.0,
+        key=jax.random.key(42)
     )
     state = dp_agg.init(self.params)
     update_fn = self.variant(dp_agg.update)
@@ -87,7 +91,9 @@ def test_clipping_norm(self, l2_norm_clip):
   def test_noise_multiplier(self, l2_norm_clip, noise_multiplier):
     """Standard dev. of noise should be l2_norm_clip * noise_multiplier."""
     dp_agg = _privacy.differentially_private_aggregate(
-        l2_norm_clip=l2_norm_clip, noise_multiplier=noise_multiplier, seed=1337
+        l2_norm_clip=l2_norm_clip,
+        noise_multiplier=noise_multiplier,
+        key=jax.random.key(1337)
     )
     state = dp_agg.init(self.params)
     update_fn = self.variant(dp_agg.update)
@@ -103,7 +109,9 @@ def test_noise_multiplier(self, l2_norm_clip, noise_multiplier):
   def test_aggregated_updates_as_input_fails(self):
     """Expect per-example gradients as input to this transform."""
     dp_agg = _privacy.differentially_private_aggregate(
-        l2_norm_clip=0.1, noise_multiplier=1.1, seed=2021
+        l2_norm_clip=0.1,
+        noise_multiplier=1.1,
+        key=jax.random.key(2021)
     )
     state = dp_agg.init(self.params)
     mean_grads = jax.tree.map(lambda g: g.mean(0), self.per_eg_grads)

optax/monte_carlo/control_variates.py

@@ -310,7 +310,7 @@ def control_variates_jacobians(
   # gradient.
   # The rng has to be the same as passed to the grad_estimator above so that we
   # obtain the same samples.
-  samples = dist_builder(*params).sample((num_samples,), seed=rng)
+  samples = dist_builder(*params).sample((num_samples,), key=rng)
   # If the CV has state, update it.
   control_variate_state = update_state_cv(
       params, samples, control_variate_state

optax/monte_carlo/stochastic_gradient_estimators.py

@@ -85,7 +85,7 @@ def score_function_jacobians(
   def surrogate(params):
     dist = dist_builder(*params)
     one_sample_surrogate_fn = lambda x: function(x) * dist.log_prob(x)
-    samples = jax.lax.stop_gradient(dist.sample((num_samples,), seed=rng))
+    samples = jax.lax.stop_gradient(dist.sample((num_samples,), key=rng))
     # We vmap the function application over samples - this ensures that the
     # function we use does not have to be vectorized itself.
     return jax.vmap(one_sample_surrogate_fn)(samples)
@@ -141,7 +141,7 @@ def surrogate(params):
     # We vmap the function application over samples - this ensures that the
     # function we use does not have to be vectorized itself.
     dist = dist_builder(*params)
-    return jax.vmap(function)(dist.sample((num_samples,), seed=rng))
+    return jax.vmap(function)(dist.sample((num_samples,), key=rng))
 
   return jax.jacfwd(surrogate)(params)
 
@@ -239,7 +239,7 @@ def measure_valued_estimation_mean(
   mean, log_std = dist.params
   std = jnp.exp(log_std)
 
-  dist_samples = dist.sample((num_samples,), seed=rng)
+  dist_samples = dist.sample((num_samples,), key=rng)
 
   pos_rng, neg_rng = jax.random.split(rng)
   pos_sample = jax.random.weibull_min(
@@ -312,7 +312,7 @@ def measure_valued_estimation_std(
   mean, log_std = dist.params
   std = jnp.exp(log_std)
 
-  dist_samples = dist.sample((num_samples,), seed=rng)
+  dist_samples = dist.sample((num_samples,), key=rng)
 
   pos_rng, neg_rng = jax.random.split(rng)
 

optax/perturbations/_make_pert.py

@@ -25,6 +25,7 @@
 import jax.numpy as jnp
 from optax import tree_utils as otu
 from optax._src import base
+from optax._src import utils
 
 
 Shape = base.Shape
@@ -35,11 +36,11 @@ class Normal:
 
   def sample(
       self,
-      seed: chex.PRNGKey,
+      key: chex.PRNGKey | int,
       sample_shape: Shape,
       dtype: chex.ArrayDType = float,
   ) -> jax.Array:
-    return jax.random.normal(seed, sample_shape, dtype)
+    return jax.random.normal(utils.to_random_key(key), sample_shape, dtype)
 
   def log_prob(self, inputs: jax.Array) -> jax.Array:
     return -0.5 * inputs**2
@@ -50,11 +51,11 @@ class Gumbel:
 
   def sample(
       self,
-      seed: chex.PRNGKey,
+      key: chex.PRNGKey | int,
       sample_shape: Shape,
       dtype: chex.ArrayDType = float,
   ) -> jax.Array:
-    return jax.random.gumbel(seed, sample_shape, dtype)
+    return jax.random.gumbel(utils.to_random_key(key), sample_shape, dtype)
 
   def log_prob(self, inputs: jax.Array) -> jax.Array:
     return -inputs - jnp.exp(-inputs)