def apply_broadcasted_dropout(vecs, dropout_rate, rng):
  """Apply dropout, broadcasted across all but the last dimension of `vecs`."""
  if dropout_rate > 0.0:
    assert rng is not None
    keep_prob = 1.0 - dropout_rate
    keep = fastmath.random.bernoulli(rng, keep_prob, (vecs.shape[-1],))
    multiplier = keep.astype(vecs.dtype) / keep_prob
    return vecs * multiplier
  else:
    return vecs
Apply dropout, broadcasted across all but the last dimension of `vecs`.
trax/layers/research/efficient_attention.py