code: def discount_reward(
      reward,
      sequence_length=None,
      discount=1.0,
      normalize=False,
      dtype=None,
      tensor_rank=1,
  ):
      """Computes discounted reward.
  
      :attr:`reward` and :attr:`sequence_length` can be either Tensors or python
      arrays. If both are python array (or `None`), the return will be a python
      array as well. Otherwise tf Tensors are returned.
  
      Args:
          reward: A Tensor or python array. Can be 1D with shape `[batch_size]`,
              or 2D with shape `[batch_size, max_time]`.
          sequence_length (optional): A Tensor or python array of shape
              `[batch_size]`. Time steps beyond the respective sequence lengths
              will be masked. Required if :attr:`reward` is 1D.
          discount (float): A scalar. The discount factor.
          normalize (bool): Whether to normalize the discounted reward, by
              `(discounted_reward - mean) / std`. Here `mean` and `std` are
              over all time steps and all samples in the batch.
          dtype (dtype): Type of :attr:`reward`. If `None`, infer from
              `reward` automatically.
          tensor_rank (int): The number of dimensions of :attr:`reward`.
              Default is 1, i.e., :attr:`reward` is a 1D Tensor consisting
              of a batch dimension. Ignored if :attr:`reward`
              and :attr:`sequence_length` are python arrays (or `None`).
  
      Returns:
          A 2D Tensor or python array of the discounted reward.
  
          If :attr:`reward` and :attr:`sequence_length` are python
          arrays (or `None`), the returned value is a python array as well.
  
  
      Example:
  
          .. code-block:: python
  
              r = [2., 1.]
              seq_length = [3, 2]
              discounted_r = discount_reward(r, seq_length, discount=0.1)
              # discounted_r == [[2. * 0.1^2, 2. * 0.1, 2.],
              #                  [1. * 0.1,   1.,       0.]]
  
              r = [[3., 4., 5.], [6., 7., 0.]]
              seq_length = [3, 2]
              discounted_r = discount_reward(r, seq_length, discount=0.1)
              # discounted_r == [[3. + 4.*0.1 + 5.*0.1^2, 4. + 5.*0.1, 5.],
              #                  [6. + 7.*0.1,            7.,          0.]]
      """
      is_tensor = tf.contrib.framework.is_tensor
      if is_tensor(reward) or is_tensor(sequence_length):
          if tensor_rank == 1:
              disc_reward = _discount_reward_tensor_1d(
                  reward, sequence_length, discount, dtype
              )
          elif tensor_rank == 2:
              disc_reward = _discount_reward_tensor_2d(
                  reward, sequence_length, discount, dtype
              )
          else:
              raise ValueError("`tensor_rank` can only be 1 or 2.")
  
          if normalize:
              mu, var = tf.nn.moments(
                  disc_reward, axes=[0, 1], keep_dims=True
              )
              disc_reward = (disc_reward - mu) / (tf.sqrt(var) + 1e-8)
      else:
          reward = np.array(reward)
          tensor_rank = reward.ndim
          if tensor_rank == 1:
              disc_reward = _discount_reward_py_1d(
                  reward, sequence_length, discount, dtype
              )
          elif tensor_rank == 2:
              disc_reward = _discount_reward_py_2d(
                  reward, sequence_length, discount, dtype
              )
          else:
              raise ValueError("`reward` can only be 1D or 2D.")
  
          if normalize:
              mu = np.mean(disc_reward)
              std = np.std(disc_reward)
              disc_reward = (disc_reward - mu) / (std + 1e-8)
  
      return disc_reward
  
imports: import numpy
  import tensorflow
  from texar.tf.utils.shapes import mask_sequences
label_desc: Use image node_id 1 for calling a global function with example usage: discount_reward(reward, sequence_length, discount, normalize, dtype, tensor_rank) and returns: disc_reward