Repository: google/jax

Issue:
Autodidax has some bugs involving reduce_sum
This works in JAX:

```python
import jax
from jax import grad, vmap
import jax.numpy as np

def simple(a):
  b = a
  t = a + b
  return t * b

def f(a):
  L0 = np.sum(simple(a))
  return L0

def g(a):
  dL0_da = vmap(grad(f), in_axes=0)(a)
  L1 = np.sum(dL0_da * dL0_da)
  return L1

key = jax.random.PRNGKey(0)
print(grad(g)(jax.random.normal(key,(2,4))))
```

However, the same code does not work in autodidax:

```
def simple(a):
  b = a
  t = a + b
  return t * b

def f(a):
  L0 = reduce_sum(simple(a))
  return L0

def g(a):
  dL0_da = vmap(grad(f), in_axes=0)(a)
  L1 = reduce_sum(dL0_da * dL0_da)
  return L1

print(grad(g)(np.random.rand(2,4)))
```

gives

```
[<ipython-input-27-03f392c346bd>](https://localhost:8080/#) in batched_f(*args)
     17     args_flat, in_tree = tree_flatten(args)
     18     in_axes_flat, in_tree2 = tree_flatten(in_axes)
---> 19     if in_tree != in_tree2: raise TypeError
     20     f_flat, out_tree = flatten_fun(f, in_tree)
     21     outs_flat = vmap_flat(f_flat, in_axes_flat, *args_flat)

TypeError:
```

Actually, I was thinking that it would fail because there doesn't appear to be any transpose rule for reduce_sum, but it doesn't seem that Autodidax got that far. What I was actually trying to do was see how Autodidax formulated the transpose rule for sum, because if you support vmap + arbitrary dims to reduce on it gets pretty complicated.


Discussion:
Thanks for finding this! (And I'm glad you're looking at Autodidax!) I'll fix it.

But in the meantime maybe take a look at [JAX's transpose rule for `reduce_sum`](https://github.com/google/jax/blob/main/jax/_src/lax/lax.py#L3560-L3566) if you haven't already. It's basically just a call to broadcast. What complexity do you have in mind?
Oh this is great! I think BroadcastInDim answers the question; in PyTorch we don't natively have this operator so we have to manually implement it with unsqueezes and expands.
On the off chance it's useful, `broadcast_in_dim` can be implemented in [a few lines of NumPy like this](https://github.com/google/jax/blob/c35a3ca0feb0cae5fca5ce3407ce38326c96d932/jax/_src/lax_reference.py#L242). We handled `np.sum` directly in the original Autograd [this way](https://github.com/HIPS/autograd/blob/01eacff7a4f12e6f7aebde7c4cb4c1c2633f217d/autograd/numpy/numpy_vjps.py#L297-L300) with a [helper](https://github.com/HIPS/autograd/blob/01eacff7a4f12e6f7aebde7c4cb4c1c2633f217d/autograd/numpy/numpy_vjps.py#L274-L285).

Resolving patch:
diff --git a/docs/autodidax.py b/docs/autodidax.py
--- a/docs/autodidax.py
+++ b/docs/autodidax.py
@@ -92,11 +92,16 @@ def mul(x, y): return bind1(mul_p, x, y)
 def neg(x): return bind1(neg_p, x)
 def sin(x): return bind1(sin_p, x)
 def cos(x): return bind1(cos_p, x)
-def reduce_sum(x, axis=None): return bind1(reduce_sum_p, x, axis=axis)
 def greater(x, y): return bind1(greater_p, x, y)
 def less(x, y): return bind1(less_p, x, y)
 def transpose(x, perm): return bind1(transpose_p, x, perm=perm)
 def broadcast(x, shape, axes): return bind1(broadcast_p, x, shape=shape, axes=axes)
+def reduce_sum(x, axis=None):
+  if axis is None:
+    axis = tuple(range(np.ndim(x)))
+  if type(axis) is int:
+    axis = (axis,)
+  return bind1(reduce_sum_p, x, axis=axis)
 
 def bind1(prim, *args, **params):
   out, = bind(prim, *args, **params)
@@ -873,8 +878,8 @@ def vectorized_unop_batching_rule(op, axis_size, vals_in, dims_in):
 
 def reduce_sum_batching_rule(axis_size, vals_in, dims_in, *, axis):
   (x,), (x_bdim,) = vals_in, dims_in
-  new_axis = axis + (x_bdim <= axis)
-  out_bdim = x_bdim - (new_axis < x_bdim)
+  new_axis = tuple(ax + (x_bdim <= ax) for ax in axis)
+  out_bdim = x_bdim - sum(ax < x_bdim for ax in axis)
   return [reduce_sum(x, new_axis)], [out_bdim]
 vmap_rules[reduce_sum_p] = reduce_sum_batching_rule
 
@@ -1271,8 +1276,10 @@ def vectorized_unop_abstract_eval(x: ShapedArray) -> List[ShapedArray]:
 abstract_eval_rules[cos_p] = vectorized_unop_abstract_eval
 abstract_eval_rules[neg_p] = vectorized_unop_abstract_eval
 
-def reduce_sum_abstract_eval(x: ShapedArray, *, axis: int) -> List[ShapedArray]:
-  new_shape = [d for i, d in enumerate(x.shape) if i != axis]
+def reduce_sum_abstract_eval(x: ShapedArray, *, axis: Tuple[int, ...]
+                             ) -> List[ShapedArray]:
+  axis_ = set(axis)
+  new_shape = [d for i, d in enumerate(x.shape) if i not in axis_]
   return [ShapedArray(tuple(new_shape), x.dtype)]
 abstract_eval_rules[reduce_sum_p] = reduce_sum_abstract_eval
 
@@ -1643,7 +1650,7 @@ def reduce_sum_translation(c, in_avals, in_vals, *, axis):
   subc = xc.XlaBuilder('add')
   shape = _xla_shape(ShapedArray((), x_aval.dtype))
   xops.Add(xops.Parameter(subc, 0, shape), xops.Parameter(subc, 1, shape))
-  return [xops.Reduce(c, [x], [zero], subc.build(), [axis])]
+  return [xops.Reduce(c, [x], [zero], subc.build(), axis)]
 xla_translations[reduce_sum_p] = reduce_sum_translation
 
 def broadcast_translation(c, in_avals, in_vals, *, shape, axes):
@@ -2201,8 +2208,9 @@ def tracers_to_jaxpr(tracers_in: List[PartialEvalTracer],
       var = constid_to_var.get(id(val))
       if var is None:
         aval = raise_to_shaped(get_aval(val))
-        var = tracer_to_var[id(t)] = constid_to_var[id(val)] = Var(aval)
+        var = constid_to_var[id(val)] = Var(aval)
         constvar_to_val[var] = val
+      tracer_to_var[id(t)] = var
     elif isinstance(t.recipe, JaxprEqnRecipe):
       if id(t.recipe) not in processed_eqns:
         eqns.append(recipe_to_eqn(tracer_to_var, t.recipe))
@@ -2556,6 +2564,11 @@ def add_transpose_rule(cts, x, y):
   return [z_bar, z_bar]
 transpose_rules[add_p] = add_transpose_rule
 
+def reduce_sum_transpose_rule(cts, x, *, axis):
+  y_bar, = cts
+  return [broadcast(y_bar, x.aval.shape, axis)]
+transpose_rules[reduce_sum_p] = reduce_sum_transpose_rule
+
 def xla_call_transpose_rule(cts, *invals, jaxpr, num_consts):
   del num_consts  # Unused
   undef_primals = [type(x) is UndefPrimal for x in invals]


Test patch: