Fix pytype error

PiperOrigin-RevId: 555254968
google · Aug 11, 2023 · 15fd356 · 15fd356
1 parent 4de699a
commit 15fd356
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/stax/elementwise_test.py b/tests/stax/elementwise_test.py
@@ -891,6 +891,7 @@ def kernel_fn_emp(x1, x2, get, params):
     kernel_fn_emp_g = jit(value_and_grad(kernel_fn_emp), static_argnames='get')
 
     def kernel_scalar_mc_grad_mean(x1, x2):
+      x1: np.ndarray
       key = random.PRNGKey(4)
       n_samples = 2**9
       k, k_grad = 0., 0.