Add missing typing.Optional type annotations to function parameters.

PiperOrigin-RevId: 376266982

Author: rchen152

spinoffs/fun_mc/fun_mc/fun_mc_lib.py

@@ -567,8 +567,8 @@ def maybe_broadcast_structure(from_structure: 'Any',
 def reparameterize_potential_fn(
     potential_fn: 'PotentialFn',
     transport_map_fn: 'TransportMap',
-    init_state: 'State' = None,
-    state_structure: 'Any' = None,
+    init_state: 'Optional[State]' = None,
+    state_structure: 'Optional[Any]' = None,
     track_volume: 'bool' = True,
 ) -> 'Tuple[PotentialFn, Optional[State]]':
   """Performs a change of variables of a potential function.
@@ -651,7 +651,7 @@ def wrapper(*args, **kwargs):
 
 def transform_log_prob_fn(log_prob_fn: 'PotentialFn',
                           bijector: 'BijectorNest',
-                          init_state: 'State' = None) -> 'Any':
+                          init_state: 'Optional[State]' = None) -> 'Any':
   """Transforms a log-prob function using a bijector.
 
   This takes a log-prob function and creates a new log-prob function that now
@@ -1041,7 +1041,7 @@ def metropolis_hastings_step(
     current_state: 'State',
     proposed_state: 'State',
     energy_change: 'FloatTensor',
-    log_uniform: 'FloatTensor' = None,
+    log_uniform: 'Optional[FloatTensor]' = None,
     seed=None) -> 'Tuple[State, MetropolisHastingsExtra]':
   """Metropolis-Hastings step.
 
@@ -1086,9 +1086,9 @@ def metropolis_hastings_step(
 
 
 @util.named_call
-def gaussian_momentum_sample(state: 'State' = None,
-                             shape: 'IntTensor' = None,
-                             dtype: 'DTypeNest' = None,
+def gaussian_momentum_sample(state: 'Optional[State]' = None,
+                             shape: 'Optional[IntTensor]' = None,
+                             dtype: 'Optional[DTypeNest]' = None,
                              seed=None) -> 'State':
   """Generates a sample from a Gaussian (Normal) momentum distribution.
 
@@ -1205,17 +1205,17 @@ def _default_hamiltonian_monte_carlo_energy_change_fn(
 def hamiltonian_monte_carlo_step(
     hmc_state: 'HamiltonianMonteCarloState',
     target_log_prob_fn: 'PotentialFn',
-    step_size: 'Any' = None,
-    num_integrator_steps: 'IntTensor' = None,
-    momentum: 'State' = None,
-    kinetic_energy_fn: 'PotentialFn' = None,
-    momentum_sample_fn: 'MomentumSampleFn' = None,
+    step_size: 'Optional[Any]' = None,
+    num_integrator_steps: 'Optional[IntTensor]' = None,
+    momentum: 'Optional[State]' = None,
+    kinetic_energy_fn: 'Optional[PotentialFn]' = None,
+    momentum_sample_fn: 'Optional[MomentumSampleFn]' = None,
     integrator_trace_fn: 'Callable[[IntegratorStepState, IntegratorStepExtras],'
     'TensorNest]' = lambda *args: (),
-    log_uniform: 'FloatTensor' = None,
+    log_uniform: 'Optional[FloatTensor]' = None,
     integrator_fn=None,
     unroll_integrator: 'bool' = False,
-    max_num_integrator_steps: 'IntTensor' = None,
+    max_num_integrator_steps: 'Optional[IntTensor]' = None,
     energy_change_fn:
     'Callable[[IntegratorState, IntegratorState, IntegratorExtras], '
     'Tuple[FloatTensor, Any]]' = _default_hamiltonian_monte_carlo_energy_change_fn,
@@ -1298,7 +1298,7 @@ def orig_target_log_prob_fn(x):
 
   if kinetic_energy_fn is None:
     kinetic_energy_fn = make_gaussian_kinetic_energy_fn(
-        len(target_log_prob.shape) if target_log_prob.shape is not None else tf
+        len(target_log_prob.shape) if target_log_prob.shape is not None else tf  # pytype: disable=attribute-error
         .rank(target_log_prob))
 
   if momentum_sample_fn is None:
@@ -1392,7 +1392,7 @@ def hamiltonian_integrator(
     integrator_trace_fn: 'Callable[[IntegratorStepState, IntegratorStepExtras],'
     'TensorNest]' = lambda *args: (),
     unroll: 'bool' = False,
-    max_num_steps: 'IntTensor' = None,
+    max_num_steps: 'Optional[IntTensor]' = None,
 ) -> 'Tuple[IntegratorState, IntegratorExtras]':
   """Intergrates a discretized set of Hamiltonian equations.
 
@@ -1719,7 +1719,7 @@ def _one_part(state, g, learning_rate):
 def gaussian_proposal(
     state: 'State',
     scale: 'FloatNest' = 1.,
-    seed: 'Any' = None) -> 'Tuple[State, Tuple[Tuple[()], float]]':
+    seed: 'Optional[Any]' = None) -> 'Tuple[State, Tuple[Tuple[()], float]]':
   """Axis-aligned gaussian random-walk proposal.
 
   Args:
@@ -1757,7 +1757,7 @@ def maximal_reflection_coupling_proposal(
     chain_ndims: 'int' = 0,
     scale: 'FloatNest' = 1,
     epsilon: 'FloatTensor' = 1e-20,
-    seed: 'Any' = None
+    seed: 'Optional[Any]' = None
 ) -> 'Tuple[State, Tuple[MaximalReflectiveCouplingProposalExtra, float]]':
   """Maximal reflection coupling proposal.
 
@@ -1900,7 +1900,7 @@ def random_walk_metropolis_step(
     rwm_state: 'RandomWalkMetropolisState',
     target_log_prob_fn: 'PotentialFn',
     proposal_fn: 'TransitionOperator',
-    log_uniform: 'FloatTensor' = None,
+    log_uniform: 'Optional[FloatTensor]' = None,
     seed=None) -> 'Tuple[RandomWalkMetropolisState, RandomWalkMetropolisExtra]':
   """Random Walk Metropolis Hastings `TransitionOperator`.
 
@@ -1992,8 +1992,8 @@ def running_variance_init(shape: 'IntTensor',
 def running_variance_step(
     state: 'RunningVarianceState',
     vec: 'FloatNest',
-    axis: 'Union[int, List[int], Tuple[int]]' = None,
-    window_size: 'IntNest' = None,
+    axis: 'Optional[Union[int, List[int], Tuple[int]]]' = None,
+    window_size: 'Optional[IntNest]' = None,
 ) -> 'Tuple[RunningVarianceState, Tuple[()]]':
   """Updates the `RunningVarianceState`.
 
@@ -2117,8 +2117,8 @@ def running_covariance_init(shape: 'IntTensor',
 def running_covariance_step(
     state: 'RunningCovarianceState',
     vec: 'FloatNest',
-    axis: 'Union[int, List[int], Tuple[int]]' = None,
-    window_size: 'IntNest' = None,
+    axis: 'Optional[Union[int, List[int], Tuple[int]]]' = None,
+    window_size: 'Optional[IntNest]' = None,
 ) -> 'Tuple[RunningCovarianceState, Tuple[()]]':
   """Updates the `RunningCovarianceState`.
 
@@ -2234,8 +2234,8 @@ def running_mean_init(shape: 'IntTensor',
 def running_mean_step(
     state: 'RunningMeanState',
     vec: 'FloatNest',
-    axis: 'Union[int, List[int], Tuple[int]]' = None,
-    window_size: 'IntNest' = None,
+    axis: 'Optional[Union[int, List[int], Tuple[int]]]' = None,
+    window_size: 'Optional[IntNest]' = None,
 ) -> 'Tuple[RunningMeanState, Tuple[()]]':
   """Updates the `RunningMeanState`.
 
@@ -2415,7 +2415,7 @@ def running_approximate_auto_covariance_init(
     max_lags: 'int',
     state_shape: 'IntTensor',
     dtype: 'DTypeNest',
-    axis: 'Union[int, List[int], Tuple[int]]' = None,
+    axis: 'Optional[Union[int, List[int], Tuple[int]]]' = None,
 ) -> 'RunningApproximateAutoCovarianceState':
   """Initializes `RunningApproximateAutoCovarianceState`.
 
@@ -2463,7 +2463,7 @@ def _shape_with_lags(shape):
 def running_approximate_auto_covariance_step(
     state: 'RunningApproximateAutoCovarianceState',
     vec: 'TensorNest',
-    axis: 'Union[int, List[int], Tuple[int]]' = None,
+    axis: 'Optional[Union[int, List[int], Tuple[int]]]' = None,
 ) -> 'Tuple[RunningApproximateAutoCovarianceState, Tuple[()]]':
   """Updates `RunningApproximateAutoCovarianceState`.
 
@@ -2568,7 +2568,7 @@ def _one_part(vec, buf, mean, auto_cov):
 
 
 def make_surrogate_loss_fn(
-    grad_fn: 'GradFn' = None,
+    grad_fn: 'Optional[GradFn]' = None,
     loss_value: 'tf.Tensor' = 0.,
 ) -> 'Any':
   """Creates a surrogate loss function with specified gradients.

spinoffs/oryx/oryx/experimental/matching/jax_rewrite.py

@@ -192,7 +192,7 @@ def f(x):
 """
 import functools
 
-from typing import Any, Callable, Dict, Iterator, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Tuple, Union
 
 import dataclasses
 import jax
@@ -380,7 +380,7 @@ def __str__(self):
 class Params(matcher.Pattern):
   """An immutable dictionary used to represent parameters of JAX primitives."""
 
-  def __init__(self, params: Dict[str, Any] = None, **kwargs: Any):
+  def __init__(self, params: Optional[Dict[str, Any]] = None, **kwargs: Any):
     """The constructor for a `Params` object.
 
     A `Params` object is an immutable dictionary, meant to encapsulate the