code: def mask_out_restarting(
      tensor: tf.Tensor, start_of_episode: tf.Tensor
  ):
      """Mask out `tensor` taken on the step that resets the environment.
  
      Args:
        tensor: a time-major 2-D `Tensor` of shape [T, B].
        start_of_episode: a 2-D `Tensor` of shape [T, B] that contains the points
          where the episode restarts.
  
      Returns:
        tensor of shape [T, B] with elements are masked out according to step_types,
        restarting weights of shape [T, B]
      """
      tensor.get_shape().assert_has_rank(2)
      start_of_episode.get_shape().assert_has_rank(2)
      weights = tf.cast(~start_of_episode, dtype=tf.float32)
      masked_tensor = tensor * weights
      return masked_tensor
  
file_path: acme/acme/agents/tf/svg0_prior/utils.py
imports: import collections
  from typing import Tuple, Optional, Dict, Iterable
  from acme import types
  from acme.tf import utils
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: mask_out_restarting(tensor, start_of_episode) and returns: masked_tensor