Merge pull request #4894 from LenaMartens:changelist/342264327

PiperOrigin-RevId: 342693683
jax-ml · Nov 16, 2020 · e56c3f4 · e56c3f4
2 parents 7fc5ad7 + 91777c3
commit e56c3f4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
@@ -5418,6 +5418,8 @@ def _select_and_gather_add_transpose(
     msg = ("VJP not implemented for select_and_gather (MaxPool) with window "
            "dilation, got window_dilation={}.")
     raise NotImplementedError(msg.format(window_dilation))
+  if type(t) is ad_util.Zero:
+    return [ad_util.Zero, None]
   has_base_dilation = any(d != 1 for d in base_dilation)
   if has_base_dilation:
     select_identity = (_get_max_identity if select_prim is ge_p