code: def SaveAndLoadModel(
      model: tf.Module, save_gradients=True
  ) -> tf.Module:
      # Roundtrip through saved model on disk.
      model_dir = os.path.join(
          absltest.get_default_test_tmpdir(), str(id(model))
      )
      tf.saved_model.save(
          model,
          model_dir,
          options=tf.saved_model.SaveOptions(
              experimental_custom_gradients=save_gradients
          ),
      )
      restored_model = tf.saved_model.load(model_dir)
      return restored_model
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import re
  import os
  from typing import Any, Callable, Optional
  from absl.testing import absltest
  from absl import logging
  import jax
  from jax import dtypes
  from jax import numpy
  from jax._src import test_util
  from jax import tree_util
  from jax.experimental import jax2tf
  from jax.experimental import export
  from jax._src import config
  from jax._src import xla_bridge
  import numpy
  import tensorflow
  from tensorflow.compiler.xla import xla_data_pb2
  from tensorflow.compiler.tf2xla.python import xla
label_desc: Use image node_id 3 for calling a global function with example usage: SaveAndLoadModel(model, save_gradients) and returns: restored_model