code: def make_snapshot(module: snt.Module):
      """Create a thin wrapper around a module to make it snapshottable."""
      # Get the input signature as long as it has been created.
      input_signature = _get_input_signature(module)
      if input_signature is None:
          raise ValueError(
              (
                  'module instance "{}" has no input_signature attribute, '
                  "which is required for snapshotting; run "
                  "create_variables to add this annotation."
              ).format(module.name)
          )
  
      # Wrap the module up in tf.function so we can process it properly.
      @tf.function
      def wrapped_module(*args, **kwargs):
          return module(*args, **kwargs)
  
      # pylint: disable=protected-access
      snapshot = Snapshot()
      snapshot._module = wrapped_module
      snapshot._variables = module.variables
      snapshot._trainable_variables = module.trainable_variables
      # pylint: disable=protected-access
  
      # Make sure the snapshot has the proper input signature.
      snapshot.__call__.get_concrete_function(*input_signature)
  
      # If we are an RNN also save the initial-state generating function.
      if isinstance(module, snt.RNNCore):
          snapshot.initial_state = tf.function(module.initial_state)
          snapshot.initial_state.get_concrete_function(
              tf.TensorSpec(shape=(), dtype=tf.int32)
          )
  
      return snapshot
  
imports: import abc
  import datetime
  import os
  import pickle
  import time
  from typing import Mapping, Optional, Union
  from absl import logging
  from acme import core
  from acme.utils import signals
  from acme.utils import paths
  import sonnet
  import tensorflow
  import tree
  from tensorflow.python.saved_model import revived_types
label_desc: Use image node_id 1 for calling a global function with example usage: make_snapshot(module) and returns: snapshot, module