Silence some pytype errors.

PiperOrigin-RevId: 520421455
google · Apr 19, 2023 · f35972d · f35972d
1 parent d00a6e0
commit f35972d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neural_tangents/_src/stax/requirements.py b/neural_tangents/_src/stax/requirements.py
@@ -778,7 +778,7 @@ def get_x_cov_mask(x):
   x2, cov2, mask2 = get_x_cov_mask(x2)
   nngp = _cov(x1, x2, diagonal_spatial, batch_axis, channel_axis)
 
-  ntk = np.zeros((), nngp.dtype) if compute_ntk else None
+  ntk = np.zeros((), nngp.dtype) if compute_ntk else None  # pytype: disable=attribute-error  # always-use-return-annotations
   is_gaussian = False
   is_reversed = False
   x1_is_x2 = utils.x1_is_x2(x1, x2, eps=eps)