code: def truncate_reshape_rows_by_num_steps(sample, num_steps):
      """Reshapes tensors in `sample` to have shape `[rows, num_steps, ...]`.
  
      This function takes a structure `sample` and for each tensor `t`, it truncates
      the tensor's outer dimension to be the highest possible multiple of
      `num_steps`.
  
      This is done by first calculating `rows = tf.shape(t[0]) // num_steps`, then
      truncating the `tensor` to shape `t_trunc = t[: (rows * num_steps), ...]`.
      For each tensor, it returns `tf.reshape(t_trunc, [rows, num_steps, ...])`.
  
      Args:
        sample: Nest of tensors.
        num_steps: Python integer.
  
      Returns:
        A next with tensors reshaped to `[rows, num_steps, ...]`.
      """
      first_elem = tf.nest.flatten(sample)[0]
      static_sequence_length = tf.compat.dimension_value(
          first_elem.shape[0]
      )
      if static_sequence_length is not None:
          num_rows = static_sequence_length // num_steps
          static_num_rows = num_rows
      else:
          num_rows = tf.shape(first_elem)[0] // num_steps
          static_num_rows = None
  
      def _truncate_and_reshape(t):
          truncated = t[: (num_rows * num_steps), ...]
          reshaped = tf.reshape(
              truncated,
              tf.concat(
                  ([num_rows, num_steps], tf.shape(t)[1:]), axis=0
              ),
          )
          reshaped.set_shape([static_num_rows, num_steps] + t.shape[1:])
          return reshaped
  
      return tf.nest.map_structure(_truncate_and_reshape, sample)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow.compat.v2
  from tf_agents.replay_buffers import replay_buffer
  from tf_agents.typing import types
  from tf_agents.utils import lazy_loader
label_desc: Use image node_id 2 for calling a global function with example usage: truncate_reshape_rows_by_num_steps(sample, num_steps) and returns: tf, reshaped