Fix broadcasting in BiasAdd.

PiperOrigin-RevId: 446645710

tf2jax/_src/ops_test.py

@@ -290,18 +290,25 @@ def pool(x):
   @parameterized.named_parameters(
       chex.params_product(
           (("NHWC", "NHWC"), ("NCHW", "NCHW")),
+          (
+              ("three_dims", (32, 16, 8)),
+              ("four_dims", (3, 32, 16, 8)),
+              ("five_dims", (2, 3, 32, 16, 8)),
+          ),
           named=True,
       ))
-  def test_bias_add(self, data_format):
+  def test_bias_add(self, data_format, input_shape):
     np.random.seed(42)
 
-    inputs = np.random.normal(size=(10, 32, 16, 8)).astype(np.float32)
-    bias = np.linspace(-1., 1., 8).astype(np.float32)
-
     if data_format == "NCHW":
-      inputs = np.transpose(inputs, [0, 3, 1, 2])
+      nchannels = input_shape[1]
+    elif data_format == "NHWC":
+      nchannels = input_shape[-1]
     else:
-      assert data_format == "NHWC"
+      raise ValueError(f"Unsupported format {data_format}")
+
+    inputs = np.random.normal(size=input_shape).astype(np.float32)
+    bias = bias = np.linspace(-1., 1., nchannels).astype(np.float32)
 
     def pool(x, b):
       return tf.raw_ops.BiasAdd(value=x, bias=b, data_format=data_format)

tf2jax/_src/tf2jax.py

@@ -305,17 +305,18 @@ def _bias_add(proto):
 
   data_format = str(proto.attr["data_format"].s, "utf-8")
   if data_format == "NHWC":
-    reduce_axis = (0, 1, 2)
+    expand_axis_fn = lambda x: [d for d in range(x.ndim) if d != x.ndim - 1]
   elif data_format == "NCHW":
-    reduce_axis = (0, 2, 3)
+    # TODO(shaobohou) this seems wrong but matches TF behaviour.
+    expand_axis_fn = lambda x: [d for d in range(x.ndim) if d != 1]
   else:
     raise ValueError(f"Found unsupported data format {data_format}.")
 
   def _func(value: jnp.ndarray, bias: jnp.ndarray) -> jnp.ndarray:
     if bias.ndim != 1:
       raise ValueError(
           f"Expected `bias` as a 1D array, found array with {bias.ndim} dims.")
-    bias = anp.expand_dims(bias, axis=reduce_axis)
+    bias = anp.expand_dims(bias, axis=expand_axis_fn(value))
     return anp.add(value, bias)
 
   return _func