code: def testMakeTimestepMaskWithPartialEpisode(self, allow_partial):
      first, mid, last = (
          ts.StepType.FIRST,
          ts.StepType.MID,
          ts.StepType.LAST,
      )
  
      next_step_types = tf.constant(
          [
              [mid, mid, last, first, mid, mid, last, first, mid, mid],
              [mid, mid, last, first, mid, mid, mid, mid, mid, last],
          ]
      )
      zeros = tf.zeros_like(next_step_types)
      next_time_step = ts.TimeStep(next_step_types, zeros, zeros, zeros)
  
      if not allow_partial:
          # Mask should be 0.0 for transition timesteps (3, 7) and for all timesteps
          #   belonging to the final, incomplete episode.
          expected_mask = [
              [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0],
              [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
          ]
      else:
          # Zeros only between episodes. Incomplete episodes are valid and not
          # zeroed out.
          expected_mask = [
              [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0],
              [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
          ]
      timestep_mask = ppo_utils.make_timestep_mask(
          next_time_step, allow_partial_episodes=allow_partial
      )
  
      timestep_mask_ = self.evaluate(timestep_mask)
      self.assertAllClose(expected_mask, timestep_mask_)
  
file_path: agents/tf_agents/agents/ppo/ppo_utils_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents.ppo import ppo_utils
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 1 for calling the PPOUtilsTest obj's underlying member method code with example usage: obj.testMakeTimestepMaskWithPartialEpisode(allow_partial) without return types