code: def check_unbatched_time_step_spec(
      time_step, time_step_spec, batch_size
  ):
      """Checks if time step conforms array spec, even if batched."""
      if batch_size is None:
          return array_spec.check_arrays_nest(time_step, time_step_spec)
  
      return array_spec.check_arrays_nest(
          time_step,
          array_spec.add_outer_dims_nest(time_step_spec, (batch_size,)),
      )
  
file_path: agents/tf_agents/bandits/environments/stationary_stochastic_structured_py_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import stationary_stochastic_structured_py_environment
  from tf_agents.policies import random_py_policy
  from tf_agents.specs import array_spec
label_desc: Use image node_id 1 for calling a global function with example usage: check_unbatched_time_step_spec(time_step, time_step_spec, batch_size) and returns: array_spec, array_spec