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: Constructor.
  
  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.
docstr_tok: ["Constructor",".","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: Callable[[PyTree], PyTree], param_vars: PyTree
  ):
      """Constructor.
  
      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__()
      self.variables = tf.nest.flatten(param_vars)
      self.trainable_variables = [
          v for v in self.variables if v.trainable
      ]
      self.__call__ = tf_graph
  
imports: from typing import Any, Callable, Sequence, Optional, Union
  from jax.experimental import jax2tf
  import tensorflow
  import tree
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)