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)
  
      if not all([spec.shape[0] == batch_size for spec in time_step]):
          return False
  
      unbatched_time_step = ts.TimeStep(
          step_type=time_step.step_type[0],
          reward=time_step.reward[0],
          discount=time_step.discount[0],
          observation=time_step.observation[0],
      )
      return array_spec.check_arrays_nest(
          unbatched_time_step, time_step_spec
      )
  
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_py_environment
  from tf_agents.policies import random_py_policy
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 2 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, False