Repository: google/jax

Issue:
Gradient shape error.
Consider the following code,

```python
def pairwise_displacement(Ra, Rb):
  return Ra - Rb

def periodic_displacement(side, dR):
  return np.mod(dR + side * f32(0.5), side) - f32(0.5) * side

def square_distance(dR):
  return np.sum(dR ** 2, axis=-1)


def distance(dR):
  return np.sqrt(square_distance(dR))
  
def periodic(side):
  def displacement_fn(Ra, Rb):
    return periodic_displacement(side, pairwise_displacement(Ra, Rb))
  return displacement_fn


def get_samples(sigma):
  particle_count = 10
  key = random.PRNGKey(0)

  N_2 = particle_count // 2
  particle_volume = N_2 * np.pi * (sigma ** 2 + 1.) / 4.
  box_size = np.sqrt(particle_volume / 0.96)
  
  displacement = periodic(box_size)

  R = box_size * random.uniform(key, (particle_count, 2), dtype=f64)
  displacement = vmap(vmap(displacement, (None, 0)), (0, None))
  return np.sum(displacement(R, R))


print('Got shape {} but expected a scalar!!'.format(
    grad(get_samples)(1.0).shape))
```
The function `grad(get_samples)(s)` should return a scalar. However, we get something of shape `[particle_count, particle_count, 2]`. It seems like the issue is most likely coming from tracing through `periodic(box_size)` since without that line the shapes work out as expected. 


Discussion:
Some more minimization:

```python
import jax.numpy as np
from jax import vmap, grad
from jax import random

f64 = np.float64
f32 = np.float32

def pairwise_displacement(Ra, Rb):
  return Ra - Rb

def periodic_displacement(side, dR):
  # return dR + side  # NOTE this works
  return np.mod(dR, side)


def periodic(side):
  def displacement_fn(Ra, Rb):
    return periodic_displacement(side, pairwise_displacement(Ra, Rb))
  return displacement_fn


def get_samples(sigma):
  particle_count = 10
  R = np.zeros((particle_count, 2), dtype=f64)
  displacement = vmap(vmap(periodic(sigma), (None, 0)), (0, None))
  return np.sum(displacement(R, R))

print('Got shape {} but expected a scalar!!'.format(
    grad(get_samples)(1.0).shape))
```

Notice the `# NOTE this works` comment. That makes me think there's a broadcasting issue with `np.mod`...

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1725,7 +1725,7 @@ def _div_transpose_rule(cotangent, x, y):
 rem_p = standard_binop([_num, _num], 'rem')
 ad.defjvp(rem_p,
           lambda g, x, y: _brcast(g, y),
-          lambda g, x, y: mul(neg(g), floor(div(x, y))))
+          lambda g, x, y: mul(_brcast(neg(g), x), floor(div(x, y))))
 
 
 def _broadcasting_select(c, which, x, y):


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1501,8 +1501,6 @@ def grad_test_spec(op, nargs, order, rng, dtypes, name=None, tol=None):
                    dtypes=[onp.float64]),
     grad_test_spec(lax.round, nargs=1, order=2, rng=jtu.rand_default(),
                    dtypes=[onp.float64]),
-    # grad_test_spec(lax.rem, nargs=2, order=2, rng=jtu.rand_default(),
-    #                dtypes=[onp.float64]),  # TODO(mattjj): enable
 
     grad_test_spec(lax.exp, nargs=1, order=2, rng=jtu.rand_small(),
                    dtypes=[onp.float64, onp.complex64]),
@@ -2287,6 +2285,22 @@ def f2(x, y):
     expected = onp.array(0.0)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  # TODO(mattjj): make this a more systematic test
+  def testRemainder(self):
+    rng = onp.random.RandomState(0)
+    x = rng.uniform(-0.9, 9, size=(3, 4))
+    y = rng.uniform(0.7, 1.9, size=(3, 1))
+    assert not set(onp.unique(x)) & set(onp.unique(y))
+    tol = 1e-1 if num_float_bits(onp.float64) == 32 else 1e-3
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+
+    rng = onp.random.RandomState(0)
+    x = rng.uniform(-0.9, 9, size=(1, 4))
+    y = rng.uniform(0.7, 1.9, size=(3, 4))
+    assert not set(onp.unique(x)) & set(onp.unique(y))
+    tol = 1e-1 if num_float_bits(onp.float64) == 32 else 1e-3
+    check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
+
 
 def all_bdims(*shapes):
   bdims = (itertools.chain([None], range(len(shape) + 1)) for shape in shapes)