code: def testMixturePolicyDynamicBatchSize(self):
      context_dim = 35
      observation_spec = tensor_spec.TensorSpec(
          [context_dim], tf.float32
      )
      time_step_spec = ts.time_step_spec(observation_spec)
      action_spec = tensor_spec.BoundedTensorSpec(
          shape=(), dtype=tf.int32, minimum=0, maximum=9, name="action"
      )
      sub_policies = [
          ConstantPolicy(action_spec, time_step_spec, i)
          for i in range(10)
      ]
      weights = [0, 0, 0.2, 0, 0, 0.3, 0, 0, 0.5, 0]
      dist = tfd.Categorical(probs=weights)
  
      policy = mixture_policy.MixturePolicy(dist, sub_policies)
      batch_size = tf.random.uniform(
          shape=(), minval=10, maxval=15, dtype=tf.int32
      )
      time_step = ts.TimeStep(
          tf.fill(
              tf.expand_dims(batch_size, axis=0),
              ts.StepType.FIRST,
              name="step_type",
          ),
          tf.zeros(shape=[batch_size], dtype=tf.float32, name="reward"),
          tf.ones(
              shape=[batch_size], dtype=tf.float32, name="discount"
          ),
          tf.reshape(
              tf.range(
                  tf.cast(batch_size * context_dim, dtype=tf.float32),
                  dtype=tf.float32,
              ),
              shape=[-1, context_dim],
              name="observation",
          ),
      )
      action_step = policy.action(time_step)
      actions, bsize = self.evaluate([action_step.action, batch_size])
      self.assertAllEqual(actions.shape, [bsize])
      self.assertAllInSet(actions, [2, 5, 8])
  
      train_step = tf.compat.v1.train.get_or_create_global_step()
      saver = policy_saver.PolicySaver(policy, train_step=train_step)
      location = os.path.join(self.get_temp_dir(), "saved_policy")
      if not tf.executing_eagerly():
          with self.cached_session():
              self.evaluate(tf.compat.v1.global_variables_initializer())
              saver.save(location)
      else:
          saver.save(location)
      loaded_policy = tf.compat.v2.saved_model.load(location)
      new_batch_size = 3
      new_time_step = ts.TimeStep(
          tf.fill(
              tf.expand_dims(new_batch_size, axis=0),
              ts.StepType.FIRST,
              name="step_type",
          ),
          tf.zeros(
              shape=[new_batch_size], dtype=tf.float32, name="reward"
          ),
          tf.ones(
              shape=[new_batch_size], dtype=tf.float32, name="discount"
          ),
          tf.reshape(
              tf.range(
                  tf.cast(
                      new_batch_size * context_dim, dtype=tf.float32
                  ),
                  dtype=tf.float32,
              ),
              shape=[-1, context_dim],
              name="observation",
          ),
      )
      new_action = self.evaluate(
          loaded_policy.action(new_time_step).action
      )
      self.assertAllEqual(new_action.shape, [new_batch_size])
      self.assertAllInSet(new_action, [2, 5, 8])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import mixture_policy
  from tf_agents.policies import policy_saver
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 3 for calling the MixturePolicyTest obj's underlying member method code with example usage: obj.testMixturePolicyDynamicBatchSize() without return types