code: def test_serialized_model_discrete(self):
      vocab_size = 3
      obs = np.array([[[0, 1], [1, 1], [1, 0], [0, 0]]])
      act = np.array([[1, 0, 0]])
      mask = np.array([[1, 1, 1, 0]])
  
      test_model_inputs = []
  
      # pylint: disable=invalid-name
      def TestModelSavingInputs(mode):
          del mode
  
          def f(inputs):
              # Save the inputs for a later check.
              test_model_inputs.append(inputs)
              # Change type to np.float32 and add the logit dimension.
              return jnp.broadcast_to(
                  inputs.astype(np.float32)[:, :, None],
                  inputs.shape + (vocab_size,),
              )
  
          return layers_base.Fn("TestModelSavingInputs", f)
          # pylint: enable=invalid-name
  
      obs_serializer = space_serializer.create(
          gym.spaces.MultiDiscrete([2, 2]), vocab_size=vocab_size
      )
      act_serializer = space_serializer.create(
          gym.spaces.Discrete(2), vocab_size=vocab_size
      )
      serialized_model = serialization_utils.SerializedModel(
          TestModelSavingInputs,  # pylint: disable=no-value-for-parameter
          observation_serializer=obs_serializer,
          action_serializer=act_serializer,
          significance_decay=0.9,
      )
  
      example = (obs, act, obs, mask)
      serialized_model.init(shapes.signature(example))
  
      (obs_logits, obs_repr, weights) = serialized_model(example)
      # Check that the model has been called with the correct input.
      np.testing.assert_array_equal(
          # The model is called multiple times for determining shapes etc.
          # Check the last saved input - that should be the actual concrete array
          # calculated during the forward pass.
          test_model_inputs[-1],
          # Should be serialized observations and actions interleaved.
          [[0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0]],
      )
      # Check the output shape.
      self.assertEqual(obs_logits.shape, obs_repr.shape + (vocab_size,))
      # Check that obs_logits are the same as obs_repr, just broadcasted over the
      # logit dimension.
      np.testing.assert_array_equal(
          np.min(obs_logits, axis=-1), obs_repr
      )
      np.testing.assert_array_equal(
          np.max(obs_logits, axis=-1), obs_repr
      )
      # Check that the observations are correct.
      np.testing.assert_array_equal(obs_repr, obs)
      # Check weights.
      np.testing.assert_array_equal(
          weights,
          [[[1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [0.0, 0.0]]],
      )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import gin
  import gym
  from jax import numpy
  import numpy
  from trax import models
  from trax import shapes
  from trax import test_utils
  from trax.data import inputs
  from trax.layers import base
  from trax.models import transformer
  from trax.rl import serialization_utils
  from trax.rl import space_serializer
  from trax.supervised import trainer_lib
label_desc: Use image node_id 2 for calling the SerializationTest obj's underlying member method code with example usage: obj.test_serialized_model_discrete() and returns: layers_base, jnp