code: def truncation(
      observation: types.NestedTensorOrArray,
      reward: types.NestedTensorOrArray,
      discount: types.Float = 1.0,
      outer_dims: Optional[types.Shape] = None,
  ) -> TimeStep:
      """Returns a `TimeStep` with `step_type` set to `StepType.LAST`.
  
      If `discount` is a scalar, and `observation` contains Tensors,
      then `discount` will be broadcasted to match the outer dimensions.
  
      Args:
        observation: A NumPy array, tensor, or a nested dict, list or tuple of
          arrays or tensors.
        reward: A NumPy array, tensor, or a nested dict, list or tuple of arrays or
          tensors.
        discount: (optional) A scalar, or 1D NumPy array, or tensor.
        outer_dims: (optional) If provided, it will be used to determine the batch
          dimensions. If not, the batch dimensions will be inferred by reward's
          shape.
  
      Returns:
        A `TimeStep`.
      """
      first_observation = tf.nest.flatten(observation)[0]
      if not tf.is_tensor(first_observation):
          if outer_dims is not None:
              step_type = np.tile(StepType.LAST, outer_dims)
              discount = _as_array(discount)
              return TimeStep(step_type, reward, discount, observation)
          # Infer the batch size.
          reward = tf.nest.map_structure(
              lambda x: _as_array(x, _get_np_dtype(x)), reward
          )
          first_reward = tf.nest.flatten(reward)[0]
          discount = _as_array(discount)
          if first_reward.shape:
              step_type = np.tile(StepType.LAST, first_reward.shape)
          else:
              step_type = StepType.LAST
          return TimeStep(step_type, reward, discount, observation)
  
      reward = tf.nest.map_structure(
          lambda r: tf.convert_to_tensor(
              value=r, dtype=r.dtype, name="reward"
          ),
          reward,
      )
      if outer_dims is not None:
          shape = outer_dims
      else:
          first_reward = tf.nest.flatten(reward)[0]
          if first_reward.shape.rank == 0:
              shape = []
          else:
              shape = [
                  tf.compat.dimension_value(first_reward.shape[0])
                  or tf.shape(input=first_reward)[0]
              ]
      step_type = tf.fill(shape, StepType.LAST, name="step_type")
      discount = tf.convert_to_tensor(
          value=discount, dtype=tf.float32, name="discount"
      )
      if discount.shape.rank == 0:
          discount = tf.fill(shape, discount, name="discount_fill")
      return TimeStep(step_type, reward, discount, observation)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import pprint
  from typing import NamedTuple, Optional
  import numpy
  import tensorflow
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
label_desc: Use image node_id 7 for calling a global function with example usage: truncation(observation, reward, discount, outer_dims) and returns: TimeStep, TimeStep, TimeStep