code: def mask_discount(discount, discount_mask):
      """Computes a discount to apply at a given timestep, based on the mask."""
      return fastmath.numpy.where(discount_mask, discount, 1.0)
  
file_path: trax/trax/rl/advantages.py
imports: import gin
  import numpy
  from trax import fastmath
label_desc: Use image node_id 1 for calling a global function with example usage: mask_discount(discount, discount_mask) and returns: fastmath