Repository: google/jax

Issue:
lax.broadcast_in_dim transposition rule
I noticed that the transposition rule of lax.broadcast_in_dim assumes that the input operand has a shape that matches the broadcast dimensions in the target shape. That is, we cannot use the following
```
x = np.ones((1, 3, 1))
lax.broadcast_in_dim(x, (2, 3, 4), broadcast_dimensions=(0, 1, 2))
```
The above would compute correctly but will fail in transpose (only if core.skip_checks == False). 

Instead one must do:
```
x = np.ones((1, 3, 1))
lax.broadcast_in_dim(x.reshape((3,)), (2, 3, 4), broadcast_dimensions=(1,))
```

Either we must add a new parameter to broadcast_in_dim with the input shape, or at the very least we ought to change the shape rule to reject unsupported uses.


Discussion:
We may be able to do this without a new parameter by accessing the input avals, similar to the approach that @mattjj recently used to drop the `input_dtype` argument from `convert_element_type_p`.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3301,9 +3301,12 @@ def _broadcast_in_dim_shape_rule(operand, *, shape, broadcast_dimensions):
 
   return shape
 
-def _broadcast_in_dim_transpose_rule(t, *, shape, broadcast_dimensions):
-  axes = tuple(np.delete(range(len(shape)), broadcast_dimensions))
-  return [_reduce_sum(t, axes)]
+def _broadcast_in_dim_transpose_rule(ct, operand, *, shape, broadcast_dimensions):
+  shape_in = operand.aval.shape
+  unit_dimensions = tuple(i for i, s in enumerate(shape_in) if s == 1)
+  bdims = tuple(np.delete(broadcast_dimensions, unit_dimensions))
+  axes = tuple(np.delete(range(len(shape)), bdims))
+  return [expand_dims(_reduce_sum(ct, axes), unit_dimensions)]
 
 def _broadcast_in_dim_batch_rule(batched_args, batch_dims, *, shape,
                                  broadcast_dimensions):
@@ -3318,7 +3321,7 @@ def _broadcast_in_dim_batch_rule(batched_args, batch_dims, *, shape,
 broadcast_in_dim_p = standard_primitive(
     _broadcast_in_dim_shape_rule, _input_dtype, 'broadcast_in_dim')
 broadcast_in_dim_p.def_impl(_broadcast_in_dim_impl)
-ad.deflinear(broadcast_in_dim_p, _broadcast_in_dim_transpose_rule)
+ad.deflinear2(broadcast_in_dim_p, _broadcast_in_dim_transpose_rule)
 batching.primitive_batchers[broadcast_in_dim_p] = _broadcast_in_dim_batch_rule
 
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1021,6 +1021,15 @@ def testBroadcastInDim(self, inshape, dtype, outshape, dimensions):
     op = lambda x: lax.broadcast_in_dim(x, outshape, dimensions)
     self._CompileAndCheck(op, args_maker)
 
+  def testBroadcastInDimOperandShapeTranspose(self):
+    # Regression test for https://github.com/google/jax/issues/5276
+    def f(x):
+      return lax.broadcast_in_dim(x, (2, 3, 4), broadcast_dimensions=(0, 1, 2)).sum()
+    def g(x):
+      return lax.broadcast_in_dim(x.reshape((3,)), (2, 3, 4), broadcast_dimensions=(1,)).sum()
+    x = np.ones((1, 3, 1))
+    self.assertArraysEqual(jax.grad(f)(x), jax.grad(g)(x))
+
   @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_inshape={}_outshape={}_bcdims={}".format(
       jtu.format_shape_dtype_string(inshape, np.float32),