Repository: google/jax

Issue:
dev timeline for more batched ops
Hi JAX Team,

I'm interested in using your library, e.g. to use gradients to optimize functions whose
terms include functions g() of jacobians of functions f(). There's no good way in PyTorch to do this for f() with output dim > 1, and certainly no good way to do this in a batched setting.

In JAX, i'm able to very easily get e.g. the trace of the jacobian of a plain relu neural net. So cool! Things break down when I try to use vmap to get a batched version of the computation.  I'm wondering, do you have a rough guess at the timeline for when I might be able to vmap in this kind of setting? Note: I've successfully used JAX to do vmap's of functions g() of jacobians of functions f() for simple choices of f()

```python
import jax.numpy as np
from jax import jit, jacrev, vmap

W = np.eye(3)
b = np.zeros(3)

def relu(x):
    return np.maximum(0,x)

def NN(x):
    return relu(np.dot(W,x) + b)

Jx_NN = jit(jacrev(NN))

# By the way, is there a recommended way to compose already JIT'ed functions?
def trace_Jx_NN(x):
    J = Jx_NN(x)
    return np.trace(J) 


x1 = np.array([2.0,2.0,2.0])
x2 = np.array([3.0,3.0,3.0])
X = np.vstack((x1,x2))

# this works
print(trace_Jx_NN(x1)) 

# this line executes
batched_trace_Jx_NN = jit(vmap(trace_Jx_NN,in_axes=(0)))

# raises "NotImplementedError  # TODO(schsam, mattjj): Handle more cases."
print(batched_trace_Jx_NN(X)) 
```

Thanks!
mark



Discussion:
Dev timeline! What kind of professional and organized operation do you think we're running here? :)

More seriously, thanks for your interest, and the kind words about JAX! To answer your question, our work is driven primarily in two ways: first, we're pushing forward the core system as a research project (e.g. with new function transformations for parallel programming and automatic masking), and second, we're extending the system's coverage, adding miscellaneous features and docs, and squashing bugs. There's an approximately infinite amount of work we could do in the latter category, and so we prioritize that work based on what users ask for.

In other words, _you_ drive our timeline, along with all the other users kind enough to open GitHub issues and provide feedback.

Even better, you can join us and help accomplish these goals by contributing code. One of the strengths of JAX is it's a pretty small pure-Python system, so once you get over the hump of learning about the codebase, it should be relatively easy for any Python programmer to contribute. (But one of the weaknesses of JAX is that we haven't yet written down enough documentation to help new developers get on board, so that initial hump is much higher than it needs to be. We're working on that!) 

On your specific question about this batching rule (as we call the primitive-by-primitive transformation rules for `vmap`), let's consider this issue a feature request for it. @sschoenholz do you have a sense for how tricky it would be to cover the case needed by this issue? I believe [this is the error being raised](https://github.com/google/jax/blob/1ab4a2ea541f83b03249e78f2c15e458d409f6e2/jax/lax.py#L1762).

If you wanted to dig in, that line is where you'd want to start adding code. Here's a quick attempt at explaining what's going on. That function, `_select_batch_rule`, is basically responsible for dealing with a `lax.select` in the context of a `vmap` transformation. The game is this: the user code called `lax.select` (maybe indirectly through a numpy function, in this case `np.maximum`), blissfully ignorant that, because of a `vmap`, there was an extra batch dimension lurking behind the scenes. But someone has to deal with that extra dimension when it comes time to evaluate the `lax.select` call, and that someone is `_select_batch_rule`. To do that, `_select_batch_rule` is given information about what it has to deal with: [the full arguments, with any batch dimensions exposed](https://github.com/google/jax/blob/1ab4a2ea541f83b03249e78f2c15e458d409f6e2/jax/lax.py#L1758), and [their corresponding batch dimensions](https://github.com/google/jax/blob/1ab4a2ea541f83b03249e78f2c15e458d409f6e2/jax/lax.py#L1759), represented as an integer (indicating which dimension/axis is the one hidden from the `lax.select` being called) or a None if that argument doesn't have a hidden dimension after all.

It looks like we've only implemented the logic for the cases where the batch dimensions corresponding to `on_true` and `on_false` are either `None` (i.e. no hidden dimension to deal with) or equal to `pred_dim`. One easy way to handle the other cases might be just to transpose those dimensions to the front; I suspect we have to do something like that anyway.
I edited the original post to add these lines to the top of the code, making it a full repro:

```python
import jax.numpy as np
from jax import jit, jacrev, vmap

W = np.eye(3)
b = np.zeros(3)
```

Please revise further if that edit was misguided.

Resolving patch:
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -265,9 +265,9 @@ def zeros_like_batched(batched_args, batch_dims):
 # method. To handle that case, the `broadcast` function uses a try/except.
 
 
-def bdim_at_front(x, bdim, broadcast_size=1):
+def bdim_at_front(x, bdim, broadcast_size=1, force_broadcast=False):
   if bdim is None:
-    return broadcast(x, broadcast_size)
+    return broadcast(x, broadcast_size, force_broadcast=force_broadcast)
   else:
     return move_dim_to_front(x, bdim)
 
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1755,17 +1755,30 @@ def _select_transpose_rule(t, pred, on_true, on_false):
             select(pred, zeros, t) if on_false is None else None]
 
 def _select_batch_rule(batched_args, batch_dims, **unused_kwargs):
-  oprand, on_true, on_false, = batched_args
+  pred, on_true, on_false, = batched_args
   pred_bdim, ot_bdim, of_bdim = batch_dims
+  size = next(x.shape[i] for x, i in zip(batched_args, batch_dims)
+              if i is not None)
 
-  if (ot_bdim not in {None, pred_bdim}) or (of_bdim not in {None, pred_bdim}):
-    raise NotImplementedError  # TODO(schsam, mattjj): Handle more cases.
-
-  # TODO(schsam, mattjj): Switch to using broadcast_in_dim.
-  ot = _ones(oprand) * on_true
-  of = _ones(oprand) * on_false
-
-  return select(oprand, ot, of), pred_bdim
+  # avoid transposes and some broadcasts in special cases
+  if pred_bdim == ot_bdim == of_bdim:
+    if onp.shape(pred) == onp.shape(on_true):
+      return select(pred, on_true, on_false), pred_bdim
+    else:
+      # vmapped function had a scalar pred with nonscalar args
+      assert onp.ndim(pred) == 1
+      pred = broadcast_in_dim(pred, on_true.shape, [pred_bdim])
+      return select(pred, on_true, on_false), pred_bdim
+
+  pred = batching.bdim_at_front(pred, pred_bdim, size, force_broadcast=True)
+  on_true = batching.bdim_at_front(on_true, ot_bdim, size, force_broadcast=True)
+  on_false = batching.bdim_at_front(on_false, of_bdim, size, force_broadcast=True)
+  assert onp.shape(on_true) == onp.shape(on_false)
+  if 0 < onp.ndim(pred) < onp.ndim(on_true):
+    # vmapped function had a scalar pred with nonscalar args
+    assert onp.ndim(pred) == 1
+    pred = broadcast_in_dim(pred, on_true.shape, [0])
+  return select(pred, on_true, on_false), 0
 
 select_p = standard_primitive(_select_shape_rule, _select_dtype_rule, 'select')
 ad.defjvp(select_p,


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -485,5 +485,43 @@ def f(params, x):
     per_example_direct = np.concatenate(per_example_direct, axis=0)
     self.assertAllClose(per_example, per_example_direct, check_dtypes=True)
 
+  def testSelect(self):
+    pred = onp.array([True, False])
+    on_true = onp.array([0, 1])
+    on_false = onp.array([2, 3])
+    ans = vmap(lax.select)(pred, on_true, on_false)
+    expected = onp.array([0, 3])
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
+    pred = onp.array([False, True])
+    on_true = onp.array([0, 1])
+    on_false = onp.array([2, 3])
+    ans = vmap(lax.select, (0, None, None))(pred, on_true, on_false)
+    expected = onp.array([[2, 3],
+                          [0, 1]])
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
+    pred = True
+    on_true = onp.array([0, 1], onp.float32)
+    on_false = onp.array(3, onp.float32)
+    ans = vmap(lax.select, (None, 0, None))(pred, on_true, on_false)
+    expected = onp.array([0, 1], onp.float32)
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
+    pred = onp.array([False, True])
+    on_true = onp.array([0, 1], onp.float32)
+    on_false = onp.array(3, onp.float32)
+    ans = vmap(lax.select, (0, 0, None))(pred, on_true, on_false)
+    expected = onp.array([3, 1], onp.float32)
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
+    pred = onp.array([False, True])
+    on_true = onp.array([2], onp.float32)
+    on_false = onp.array([[3, 4]], onp.float32)
+    ans = vmap(lax.select, (0, None, 1), 1)(pred, on_true, on_false)
+    expected = onp.array([[3, 2]], onp.float32)
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
+
 if __name__ == '__main__':
   absltest.main()