def SplitControls():  # pylint: disable=invalid-name
    """Splits logits for actions in different controls."""
    def f(x):
      return jnp.reshape(x, x.shape[:2] + (n_controls, n_actions))
    return tl.Fn('SplitControls', f)
Splits logits for actions in different controls.
trax/rl/serialization_utils.py