Gradient through Neural Tangent Kernel works for CNN without MaxPool, but not for CNN with MaxPool
Hi all,

I am working on a project that requires taking the gradient of Neural Tangent Kernel with respect to the parameters at which the Jacobian is being evaluated.

However, I’m getting an error when jitting the NTK computation function if my network uses maxpooling.

I have prepared two Colab notebooks that reproduce the issue:
[CNN without MaxPool](https://colab.research.google.com/drive/1k_cCuNxz2rXeXFrVnwFimWtvOzfVPjdo?usp=sharing)
[CNN with MaxPool](https://colab.research.google.com/drive/1WxK6WjSrZ4FdzJpY3hg0LSZ8gSkM7Bf-?usp=sharing)

In the first notebook, I show that computing the gradients with respect to the NTK evaluation parameters works just fine if I don’t include a maxpooling operation in the CNN, and in the second notebook, I show that including a maxpooling operation results in an error:
> typeValue Zero(ShapedArray(float32[11,28,28,32])) with type <class ‘jax.ad_util.Zero’> is not a valid JAX type

Oddly, the error message disappears if I remove all the jitting from the NTK function.

I would really appreciate if you could help me understand the cause of this issue and help me fix it.

Thank you so much!
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index f8944b13a..e2df6b7b0 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5573,6 +5573,8 @@ def _select_and_scatter_add_transpose(
     t, source, operand, *, select_prim, window_dimensions, window_strides,
     padding):
   assert ad.is_undefined_primal(source) and not ad.is_undefined_primal(operand)
+  if type(t) is ad_util.Zero:
+    return [ad_util.Zero(source.aval), None]
   ones = (1,) * len(window_dimensions)
   source_t = _select_and_gather_add(t, operand, select_prim, window_dimensions,
                                     window_strides, padding, ones, ones)
['jax/_src/lax/lax.py']