code: def make_sequences_from_transitions(
      transitions: types.Transition, num_batch_dims: int = 1
  ) -> adders.Step:
      """Convert a batch of transitions into a batch of 1-step sequences."""
      stack = lambda x, y: jnp.stack((x, y), axis=num_batch_dims)
      duplicate = lambda x: stack(x, x)
      observation = jax.tree_map(
          stack, transitions.observation, transitions.next_observation
      )
      reward = duplicate(transitions.reward)
  
      return adders.Step(  # pytype: disable=wrong-arg-types  # jnp-type
          observation=observation,
          action=duplicate(transitions.action),
          reward=reward,
          discount=duplicate(transitions.discount),
          start_of_episode=jnp.zeros_like(reward, dtype=jnp.bool_),
          extras=jax.tree_map(duplicate, transitions.extras),
      )
  
file_path: acme/acme/agents/jax/mpo/utils.py
imports: from typing import Callable
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import types
  import distrax
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 5 for calling a global function with example usage: make_sequences_from_transitions(transitions, num_batch_dims) and returns: adders