import tensorflow as tf

      to_seq = mask.shape.dims + [1] * (sequence.shape.ndims - mask.shape.ndims)
      sequence *= tf.reshape(mask, to_seq)
      to_decay = mask.shape.dims + [1] * (decay.shape.ndims - mask.shape.ndims)
      decay *= tf.reshape(mask, to_decay)

    sequences = [sequence, decay]
