Trimming the library.

Remove tests of deprecated functions.
Reschedule deletion of deprecated functions for earlier future releases.

PiperOrigin-RevId: 777652709

Author: vroulet

docs/api/control_variates.rst

@@ -1,16 +1,32 @@
-Stochastic Gradient Estimators
-==============================
+Stochastic Gradient Estimators and Control Variates
+===================================================
 
 .. warning::
-    This module has been deprecated and will be removed in optax 0.3.0.
+    This module has been deprecated and will be removed in optax 0.2.7
 
 .. currentmodule:: optax.monte_carlo
 
 .. autosummary::
+    control_delta_method
+    control_variates_jacobians
     measure_valued_jacobians
+    moving_avg_baseline
     pathwise_jacobians
     score_function_jacobians
 
+
+Control delta method
+~~~~~~~~~~~~~~~~~~~~
+.. autofunction:: control_delta_method
+
+Control variates Jacobians
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. autofunction:: control_variates_jacobians
+
+Moving average baseline
+~~~~~~~~~~~~~~~~~~~~~~~
+.. autofunction:: moving_avg_baseline
+
 Measure valued Jacobians
 ~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: measure_valued_jacobians

docs/api/stochastic_gradient_estimators.rst

@@ -1358,14 +1358,14 @@ def noisy_sgd(
   """
   if seed is not None:
     warnings.warn(
-        '"seed" is deprecated and will be removed in optax 0.3.0, use "key".',
+        '"seed" is deprecated and will be removed in optax 0.2.7, use "key".',
         DeprecationWarning,
     )
     if key is not None:
       raise ValueError('Only one of seed or key can be specified.')
     key = jax.random.key(seed)
   if key is None:
-    warnings.warn('Specifying a key will be required in optax 0.3.0.')
+    warnings.warn('Specifying a key will be required in optax 0.2.7.')
     key = jax.random.key(0)
   key = utils.canonicalize_key(key)
 

optax/_src/alias.py

@@ -70,14 +70,14 @@ def differentially_private_aggregate(
 
   if seed is not None:
     warnings.warn(
-        '"seed" is deprecated and will be removed in optax 0.3.0, use "key".',
+        '"seed" is deprecated and will be removed in optax 0.2.7, use "key".',
         DeprecationWarning,
     )
     if key is not None:
       raise ValueError('Only one of seed or key can be specified.')
     key = jax.random.key(seed)
   if key is None:
-    warnings.warn('Specifying a key will be required in optax 0.3.0.')
+    warnings.warn('Specifying a key will be required in optax 0.2.7.')
     key = jax.random.key(0)
   key = utils.canonicalize_key(key)
 

optax/contrib/_privacy.py

@@ -93,7 +93,7 @@ def control_delta_method(
     state.
 
   .. deprecated:: 0.2.4
-    This function will be removed in 0.3.0
+    This function will be removed in 0.2.7
   """
 
   def delta(
@@ -171,7 +171,7 @@ def moving_avg_baseline(
     state.
 
   .. deprecated:: 0.2.4
-    This function will be removed in 0.3.0
+    This function will be removed in 0.2.7
   """
 
   def moving_avg(
@@ -281,7 +281,7 @@ def control_variates_jacobians(
     * The updated CV state.
 
   .. deprecated:: 0.2.4
-    This function will be removed in 0.3.0
+    This function will be removed in 0.2.7
   """
   control_variate = control_variate_from_function(function)
   stochastic_cv, expected_value_cv, update_state_cv = control_variate
@@ -418,7 +418,7 @@ def estimate_control_variate_coefficients(
     in `params`.
 
   .. deprecated:: 0.2.4
-    This function will be removed in 0.3.0
+    This function will be removed in 0.2.7
   """
   # Resample to avoid biased gradients.
   cv_rng, _ = jax.random.split(rng)