def _GateAndProject():
  """Returns a combined gating+projection layer that saves on memory."""

  def f(projection, gating, x):
    # Args arrive in reverse order because of how they were put on the stack.
    # Einsum indices: d (d_in), n (n_modules), m (d_module = d_out/n_modules)
    return jnp.einsum('...d,nd,dm->...nm', x, gating, projection)

  return tl.Fn('_GateAndProject', f)
Returns a combined gating+projection layer that saves on memory.
trax/layers/research/sparsity.py