Returns the norm for each state shard.

def _norm_squared(args: Dict[str, Any]) -> float:
    """Returns the norm for each state shard."""
    state = _state_shard(args)
    return np.sum(np.abs(state) ** 2)