def SignificanceWeights(serializer, decay):
  """Multiplies a binary mask with a symbol significance mask."""
  def significance_weights(mask):
    # (repr,) -> (batch, length, repr)
    # significance = [0, 1, 2]
    significance = serializer.significance_map
    assert significance.shape[0] == mask.shape[2]
    # significance = batch_size * [0, 1, 2]
    significance = jnp.repeat(
        significance[np.newaxis, ...], repeats=mask.shape[0], axis=0)
    # significance = batch_size * [0, 1, 2] * mask.shape[1]
    significance = jnp.repeat(
        significance[..., jnp.newaxis], repeats=mask.shape[1], axis=2)
    # significance = batch_size *  mask.shape[1] * [0, 1, 2]
    significance = jnp.swapaxes(significance, 1, 2)
    assert significance.shape == mask.shape
    sig_weights = mask * decay ** significance
    return sig_weights
  return tl.Fn('SignificanceWeights', significance_weights)
Multiplies a binary mask with a symbol significance mask.
trax/rl/serialization_utils.py