Fix a pytype error that slipped past the cleanup resulting this LSC: …

…go/lsc-jnp-ndarray-types

This is breaking upstream builds that depend on neural_tangents, such as         //learning/deepmind/public/tools/ml_python:core_deps

PiperOrigin-RevId: 513279197

neural_tangents/_src/stax/linear.py

@@ -3732,7 +3732,7 @@ def apply_fun(params, inputs, **kwargs):
 
 
 def _pos_emb_identity(shape: Sequence[int]) -> np.ndarray:
-  size = utils.size_at(shape)
+  size = utils.size_at(shape)  # pytype: disable=wrong-arg-types  # jax-ndarray
   R = np.eye(size).reshape(tuple(shape) * 2)
   R = utils.zip_axes(R)
   return R