Make conv2d tag graph matcher more general

PiperOrigin-RevId: 696493059
google-deepmind · Nov 14, 2024 · a621e70 · a621e70
1 parent e302ee3
commit a621e70
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/kfac_jax/_src/tag_graph_matcher.py b/kfac_jax/_src/tag_graph_matcher.py
@@ -1261,8 +1261,9 @@ def _conv2d(x: Array, params: Sequence[Array]) -> Array:
     # No bias
     return y
 
+  bias = params[1]
   # Add bias
-  return y + params[1][None, None, None]
+  return y + bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape)
 
 
 def _conv2d_parameter_extractor(