code: def set_expected_shape(experience, num_steps):
      """Sets expected shape."""
  
      def set_time_dim(input_tensor, steps):
          tensor_shape = input_tensor.shape.as_list()
          if len(tensor_shape) < 2:
              raise ValueError(
                  "input_tensor is expected to be of rank-2, but found otherwise: "
                  f"input_tensor={input_tensor}, tensor_shape={tensor_shape}"
              )
          tensor_shape[1] = steps
          input_tensor.set_shape(tensor_shape)
  
      tf.nest.map_structure(
          lambda t: set_time_dim(t, num_steps), experience
      )
  
file_path: agents/tf_agents/bandits/agents/examples/v2/trainer.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import logging
  import tensorflow
  from tf_agents.bandits.replay_buffers import bandit_replay_buffer
  from tf_agents.drivers import dynamic_step_driver
  from tf_agents.eval import metric_utils
  from tf_agents.metrics import export_utils
  from tf_agents.metrics import tf_metrics
  from tf_agents.policies import policy_saver
label_desc: Use image node_id 2 for calling a global function with example usage: set_expected_shape(experience, num_steps) without return types