def export_model(tokenizer_fn: Callable[[PyTree, PyTree], PyTree],
                 params: train_utils.PyTree,
                 export_dir: str,
                 with_gradient: bool,
                 input_shape: Tuple[int, ...] = (-1, 16, 128, 128, 3),
                 batch_dim: Optional[int] = None,
                 dtype=tf.float32):
  """Exports Flax model to Tensorflow Saved Model."""
  if batch_dim is None:
    polymorphic_shapes = '(batch, ...)'
  else:
    polymorphic_shapes = None
  input_shape = [batch_dim, *input_shape[1:]]
  input_signature = [tf.TensorSpec(input_shape, dtype)]

  logging.info('--> Exporting model to %s', export_dir)
  logging.info('--> Input shape %s', input_shape)

  export_utils.convert_and_save_model(
      tokenizer_fn,
      params,
      input_signatures=input_signature,
      polymorphic_shapes=polymorphic_shapes,
      model_dir=export_dir,
      with_gradient=with_gradient,
      enable_xla=True,
      compile_model=True)
  logging.info('Completed exporting model to %s', export_dir)
Exports Flax model to Tensorflow Saved Model.
videogvt/interfaces/mm_vq.py