class_docstr: Wraps a function and its parameters for saving to a SavedModel.
  
  Implements the interface described at
  https://www.tensorflow.org/hub/reusable_saved_models.
docstr: Args:
  
  tf_graph: a tf.function taking one argument (the inputs), which can be
     be tuples/lists/dictionaries of np.ndarray or tensors. The function
     may have references to the tf.Variables in `param_vars`.
  param_vars: the parameters, as tuples/lists/dictionaries of tf.Variable,
     to be saved as the variables of the SavedModel.
code: def __init__(self, tf_graph, param_vars):
      """Args:
  
      tf_graph: a tf.function taking one argument (the inputs), which can be
         be tuples/lists/dictionaries of np.ndarray or tensors. The function
         may have references to the tf.Variables in `param_vars`.
      param_vars: the parameters, as tuples/lists/dictionaries of tf.Variable,
         to be saved as the variables of the SavedModel.
      """
      super().__init__()
      # Implement the interface from https://www.tensorflow.org/hub/reusable_saved_models
      self.variables = tf.nest.flatten(param_vars)
      self.trainable_variables = [
          v for v in self.variables if v.trainable
      ]
      # If you intend to prescribe regularization terms for users of the model,
      # add them as @tf.functions with no inputs to this list. Else drop this.
      self.regularization_losses = []
      self.__call__ = tf_graph
  
file_path: jax/jax/experimental/jax2tf/examples/saved_model_lib.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  from typing import Any, Callable, Optional, Union
  from jax.experimental import jax2tf
  import tensorflow
label_desc: Use image node_id 1 to create a new _ReusableSavedModelWrapper object from inherited base classes: tf.train with example: obj = _ReusableSavedModelWrapper(tf_graph, param_vars)