<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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)