code: def test_split_batch(batch_data):
      s = TimeStepBatch(
          env_spec=batch_data["env_spec"],
          observations=batch_data["observations"],
          actions=batch_data["actions"],
          rewards=batch_data["rewards"],
          next_observations=batch_data["next_observations"],
          step_types=batch_data["step_types"],
          env_infos=batch_data["env_infos"],
          agent_infos=batch_data["agent_infos"],
          episode_infos=batch_data["episode_infos"],
      )
      batches = s.split()
  
      assert len(batches) == 2  # original batch_data is a batch of 2
      for i, batch in enumerate(batches):
          assert batch.env_spec == batch_data["env_spec"]
          assert np.array_equal(
              batch.observations, [batch_data["observations"][i]]
          )
          assert np.array_equal(
              batch.next_observations,
              [batch_data["next_observations"][i]],
          )
          assert np.array_equal(
              batch.actions, [batch_data["actions"][i]]
          )
          assert np.array_equal(
              batch.rewards, [batch_data["rewards"][i]]
          )
          assert np.array_equal(
              batch.step_types, [batch_data["step_types"][i]]
          )
          for key in batch.env_infos:
              assert key in batch_data["env_infos"]
              assert np.array_equal(
                  batch.env_infos[key],
                  [batch_data["env_infos"][key][i]],
              )
          for key in batch.agent_infos:
              assert key in batch_data["agent_infos"]
              assert np.array_equal(
                  batch.agent_infos[key],
                  [batch_data["agent_infos"][key][i]],
              )
          for key in batch.episode_infos:
              assert key in batch_data["episode_infos"]
              assert np.array_equal(
                  batch.episode_infos[key],
                  [batch_data["episode_infos"][key][i]],
              )
  
imports: import akro
  import gym.spaces
  import numpy
  import pytest
  from garage import EnvSpec, EnvStep, EpisodeBatch, StepType, TimeStep, TimeStepBatch
label_desc: Use image node_id 48 for calling a global function with example usage: test_split_batch(batch_data) without return types