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.
  
  This class contains all the code needed to convert a Flax model to a
  TensorFlow saved model.
docstr: Builds the tf.Module.
  
  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: ["Builds","the","tf.Module",".","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[[_InputBatch], _Array],
      param_vars: Dict[str, tf.Variable],
  ):
      """Builds the tf.Module.
  
      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
      ]
      self._tf_graph = tf_graph
  
imports: import functools
  from typing import Callable, Dict, List
  import absl
  from flax import linen
  from flax.metrics import tensorboard
  import jax
  from jax import numpy
  from jax.experimental import jax2tf
  import numpy
  import optax
  import tensorflow
  import tensorflow_transform
  from tfx import v1
  from tfx.examples.penguin import penguin_utils_base
label_desc: Use image node_id 1 to create a new _SavedModelWrapper object from inherited base classes: tf.train with example: obj = _SavedModelWrapper(tf_graph, param_vars)