code: """Utility functions used by PyTorch algorithms."""
  import torch
  import torch.nn.functional as F
  
  
  class _Default:  # pylint: disable=too-few-public-methods
      """A wrapper class to represent default arguments.
  
      Args:
          val (object): Argument value.
  
      """
  
      def __init__(self, val):
          self.val = val
  
  
  def make_optimizer(optimizer_type, module, **kwargs):
      """Create an optimizer for PyTorch algos.
  
      Args:
          optimizer_type (Union[type, tuple[type, dict]]): Type of optimizer.
              This can be an optimizer type such as 'torch.optim.Adam' or a
              tuple of type and dictionary, where dictionary contains arguments
              to initialize the optimizer e.g. (torch.optim.Adam, {'lr' = 1e-3})
          module (torch.nn.Module): The module whose parameters needs to be
              optimized.
          kwargs (dict): Other keyword arguments to initialize optimizer. This
              is not used when `optimizer_type` is tuple.
  
      Returns:
          torch.optim.Optimizer: Constructed optimizer.
  
      Raises:
          ValueError: Raises value error when `optimizer_type` is tuple, and
              non-default argument is passed in `kwargs`.
  
      """
      if isinstance(optimizer_type, tuple):
          opt_type, opt_args = optimizer_type
          for name, arg in kwargs.items():
              if not isinstance(arg, _Default):
                  raise ValueError('Should not specify {} and explicit \
                      optimizer args at the same time'.format(name))
          return opt_type(module.parameters(), **opt_args)
  
      opt_args = {}
      for name, arg in kwargs.items():
          if isinstance(arg, _Default):
              opt_args[name] = arg.val
          else:
              opt_args[name] = arg
      return optimizer_type(module.parameters(), **opt_args)
  
  
  def compute_advantages(discount, gae_lambda, max_path_length, baselines,
                         rewards):
      """Calculate advantages.
  
      Advantages are a discounted cumulative sum.
  
      Calculate advantages using a baseline (value function) according to
      Generalized Advantage Estimation (GAE)
  
      The discounted cumulative sum can be computed using conv2d with filter.
      filter:
          [1, (discount * gae_lambda), (discount * gae_lambda) ^ 2, ...]
          where the length is same with max_path_length.
  
      baselines and rewards are also has same shape.
          baselines:
          [ [b_11, b_12, b_13, ... b_1n],
            [b_21, b_22, b_23, ... b_2n],
            ...
            [b_m1, b_m2, b_m3, ... b_mn] ]
          rewards:
          [ [r_11, r_12, r_13, ... r_1n],
            [r_21, r_22, r_23, ... r_2n],
            ...
            [r_m1, r_m2, r_m3, ... r_mn] ]
  
      Args:
          discount (float): RL discount factor (i.e. gamma).
          gae_lambda (float): Lambda, as used for Generalized Advantage
              Estimation (GAE).
          max_path_length (int): Maximum length of a single rollout.
          baselines (torch.Tensor): A 2D vector of value function estimates with
              shape (N, T), where N is the batch dimension (number of episodes)
              and T is the maximum path length experienced by the agent. If an
              episode terminates in fewer than T time steps, the remaining
              elements in that episode should be set to 0.
          rewards (torch.Tensor): A 2D vector of per-step rewards with shape
              (N, T), where N is the batch dimension (number of episodes) and T
              is the maximum path length experienced by the agent. If an episode
              terminates in fewer than T time steps, the remaining elements in
              that episode should be set to 0.
  
      Returns:
          torch.Tensor: A 2D vector of calculated advantage values with shape
              (N, T), where N is the batch dimension (number of episodes) and T
              is the maximum path length experienced by the agent. If an episode
              terminates in fewer than T time steps, the remaining values in that
              episode should be set to 0.
  
      """
      adv_filter = torch.full((1, 1, 1, max_path_length - 1),
                              discount * gae_lambda)
      adv_filter = torch.cumprod(F.pad(adv_filter, (1, 0), value=1), dim=-1)
  
      deltas = (rewards + discount * F.pad(baselines, (0, 1))[:, 1:] - baselines)
      deltas = F.pad(deltas, (0, max_path_length - 1)).unsqueeze(0).unsqueeze(0)
  
      advantages = F.conv2d(deltas, adv_filter, stride=1).squeeze()
      return advantages
  
  
  def pad_to_last(nums, total_length, axis=-1, val=0):
      """Pad val to last in nums in given axis.
  
      length of the result in given axis should be total_length.
  
      Raises:
        IndexError: If the input axis value is out of range of the nums array
  
      Args:
          nums (numpy.ndarray): The array to pad.
          total_length (int): The final width of the Array.
          axis (int): Axis along which a sum is performed.
          val (int): The value to set the padded value.
  
      Returns:
          torch.Tensor: Padded array
  
      """
      tensor = torch.Tensor(nums)
      axis = (axis + len(tensor.shape)) if axis < 0 else axis
  
      if len(tensor.shape) <= axis:
          raise IndexError('axis {} is out of range {}'.format(
              axis, tensor.shape))
  
      padding_config = [0, 0] * len(tensor.shape)
      padding_idx = abs(axis - len(tensor.shape)) * 2 - 1
      padding_config[padding_idx] = max(total_length - tensor.shape[axis], val)
      return F.pad(tensor, padding_config)
  
  
  def filter_valids(tensor, valids):
      """Filter out tensor using valids (last index of valid tensors).
  
      valids contains last indices of each rows.
  
      Args:
          tensor (torch.Tensor): The tensor to filter
          valids (list[int]): Array of length of the valid values
  
      Returns:
          torch.Tensor: Filtered Tensor
  
      """
      return [tensor[i][:valids[i]] for i in range(len(valids))]
  
apis: torch.nn.functional.conv2d torch.nn.functional.pad torch.full torch.Tensor