code: def make_timestep_mask(
      batched_next_time_step: ts.TimeStep,
      allow_partial_episodes: bool = False,
  ) -> types.Tensor:
      """Create a mask for transitions and optionally final incomplete episodes.
  
      Args:
        batched_next_time_step: Next timestep, doubly-batched [batch_dim, time_dim,
          ...].
        allow_partial_episodes: If true, then steps on incomplete episodes are
          allowed.
  
      Returns:
        A mask, type tf.float32, that is 0.0 for all between-episode timesteps
          (batched_next_time_step is FIRST). If allow_partial_episodes is set to
          False, the mask has 0.0 for incomplete episode at the end of the sequence.
      """
      if allow_partial_episodes:
          episode_is_complete = None
      else:
          # 1.0 for timesteps of all complete episodes. 0.0 for incomplete episode at
          #   the end of the sequence.
          episode_is_complete = (
              tf.cumsum(
                  tf.cast(batched_next_time_step.is_last(), tf.float32),
                  axis=1,
                  reverse=True,
              )
              > 0
          )
  
      # 1.0 for all valid timesteps. 0.0 where between episodes.
      not_between_episodes = ~batched_next_time_step.is_first()
  
      if allow_partial_episodes:
          return tf.cast(not_between_episodes, tf.float32)
      else:
          return tf.cast(
              episode_is_complete & not_between_episodes, tf.float32
          )
  
file_path: agents/tf_agents/agents/ppo/ppo_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Sequence
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import utils
  from tf_agents.specs import distribution_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 2 for calling a global function with example usage: make_timestep_mask(batched_next_time_step, allow_partial_episodes) and returns: tf, tf