Repository: jax-ml/jax
Pull request: Implement gamma sampler using core.Primitive interface
This PR tries to address #1789. While working on a fix, I realize that (compiling time) `lax.map` outperforms `vmap` for most usage cases. So I believe that the tests for gamma/beta/dirichlet will be much faster than before. With this, probably #552 and #1188 is not necessary to be addressed anymore.

However, I am getting trouble with translation rule. With the script
```
import jax; jax.config.update('jax_platform_name', 'cpu')
from jax import vmap, random, jit, grad
import jax.numpy as np

def f(a):
    return random.gamma(random.PRNGKey(0), a)
```
`jit(f)(np.ones(3))` throws an error at printing function

<details>

```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~/miniconda3/envs/pydata/lib/python3.6/site-packages/IPython/core/formatters.py in __call__(self, obj)
    700                 type_pprinters=self.type_printers,
    701                 deferred_pprinters=self.deferred_printers)
--> 702             printer.pretty(obj)
    703             printer.flush()
    704             return stream.getvalue()

~/miniconda3/envs/pydata/lib/python3.6/site-packages/IPython/lib/pretty.py in pretty(self, obj)
    400                         if cls is not object \
    401                                 and callable(cls.__dict__.get('__repr__')):
--> 402                             return _repr_pprint(obj, self, cycle)
    403 
    404             return _default_pprint(obj, self, cycle)

~/miniconda3/envs/pydata/lib/python3.6/site-packages/IPython/lib/pretty.py in _repr_pprint(obj, p, cycle)
    695     """A pprint that just redirects to the normal repr function."""
    696     # Find newlines and replace them with p.break_()
--> 697     output = repr(obj)
    698     for idx,output_line in enumerate(output.splitlines()):
    699         if idx:

~/jax/jax/interpreters/xla.py in __repr__(self)
    655     line_width = onp.get_printoptions()['linewidth']
    656     prefix = '{}('.format(self.__class__.__name__)
--> 657     s = onp.array2string(self._value, prefix=prefix, suffix=',',
    658                          separator=', ', max_line_width=line_width)
    659     dtype_str = 'dtype={})'.format(self.dtype.name)

~/jax/jax/interpreters/xla.py in _value(self)
    608     if self._npy_value is None:
    609       self._npy_value = self.device_buffer.to_py()
--> 610       self._npy_value.flags.writeable = False
    611     return self._npy_value
    612 

AttributeError: 'tuple' object has no attribute 'flags'
```
</details>


while `jit(f)(np.ones(()))` throws the error

<details>

```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/jax/jax/abstract_arrays.py in __len__(self)
    137     try:
--> 138       return self.shape[0]
    139     except IndexError:

IndexError: tuple index out of range

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-3-6929b861a3fc> in <module>
----> 1 jit(f)(np.ones(()))

~/jax/jax/api.py in f_jitted(*args, **kwargs)
    148     _check_args(args_flat)
    149     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 150     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend)
    151     return tree_unflatten(out_tree(), out)
    152 

~/jax/jax/core.py in call_bind(primitive, f, *args, **params)
    590   if top_trace is None:
    591     with new_sublevel():
--> 592       outs = primitive.impl(f, *args, **params)
    593   else:
    594     tracers = map(top_trace.full_raise, args)

~/jax/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    373   device = params['device']
    374   backend = params.get('backend', None)
--> 375   compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
    376   try:
    377     return compiled_fun(*args)

~/jax/jax/linear_util.py in memoized_fun(fun, *args)
    207       fun.populate_stores(stores)
    208     else:
--> 209       ans = call(fun, *args)
    210       cache[key] = (ans, fun.stores)
    211     return ans

~/jax/jax/interpreters/xla.py in _xla_callable(fun, device, backend, *abstract_args)
    411   xla_consts = _map(c.Constant, consts)
    412   xla_args = _xla_callable_args(c, abstract_args, tuple_args)
--> 413   out_nodes = jaxpr_subcomp(c, jaxpr, backend, axis_env, xla_consts, (), *xla_args)
    414   built = c.Build(c.Tuple(*out_nodes))
    415 

~/jax/jax/interpreters/xla.py in jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, freevars, *args)
    260       ans = rule(c, *in_nodes, **eqn.params)
    261     elif eqn.primitive in translations:
--> 262       ans = translations[eqn.primitive](c, *in_nodes, **eqn.params)
    263     elif eqn.primitive in reduction_translations:
    264       new_params = check_backend_params(eqn.params, backend)

~/jax/jax/interpreters/xla.py in f(c, *args, **params)
    533     pvals = [pe.PartialVal((a, core.unit)) for a in avals]
    534     jaxpr, _, consts = pe.trace_to_jaxpr(
--> 535         lu.wrap_init(fun, params), pvals, instantiate=True)
    536     consts = _map(c.Constant, consts)
    537     outs = jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, (), *xla_args)

~/jax/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, **kwargs)
    313   with new_master(JaxprTrace) as master:
    314     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 315     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    316     assert not env
    317     del master

~/jax/jax/linear_util.py in call_wrapped(***failed resolving arguments***)
    155     while stack:
    156       gen, out_store = stack.pop()
--> 157       ans = gen.send(ans)
    158       if out_store is not None:
    159         ans, side = ans

~/jax/jax/interpreters/partial_eval.py in trace_to_subjaxpr(master, instantiate, pvals)
    325   in_tracers = map(trace.new_arg, pvals)
    326   ans = yield in_tracers, {}
--> 327   instantiate = [instantiate] * len(ans) if type(instantiate) is bool else instantiate
    328   out_tracers = map(trace.full_raise, map(core.full_lower, ans))
    329   out_tracers = map(partial(instantiate_const_at, trace), instantiate, out_tracers)

~/jax/jax/core.py in __len__(self)
    292 
    293   def __len__(self):
--> 294     return self.aval._len(self)
    295 
    296   @property

~/jax/jax/abstract_arrays.py in _len(self, ignored_tracer)
    141 
    142   def _len(self, ignored_tracer):
--> 143     return len(self)
    144 
    145   def strip_weak_type(self):

~/jax/jax/abstract_arrays.py in __len__(self)
    138       return self.shape[0]
    139     except IndexError:
--> 140       raise TypeError("len() of unsized object")  # same as numpy error
    141 
    142   def _len(self, ignored_tracer):

TypeError: len() of unsized object
```
</details>

I guess the output with translation rule is a tuple and I need to pack it into an array, but I am not sure how to process. I tried to mimic `custom_transform` implementation to use tree_flatten, flatten_fun_nokwargs... but don't know how to write a compatible batching rule. Several months ago, the following code works
```
    jaxpr, out, consts = partial_eval.trace_unwrapped_to_jaxpr(_gamma_impl,
                                                               tuple(lax._abstractify(o) for o in (key, a)))
    aval, _ = out
    return random_gamma_p.bind(key, a, jaxpr=jaxpr, aval=aval, consts=consts)

def _random_gamma_translate(c, key, a, jaxpr, aval, consts):
    xla_computation = xla.jaxpr_computation(jaxpr, consts, (), c.GetShape(key), c.GetShape(alpha))
    return c.Call(xla_computation, (key, alpha))
```
but those utilities `jaxpr_computation`, `trace_unwrapped_to_jaxpr` are no longer available and I am not sure if a similar code will work now.

I appreciate any help!

cc @mattjj @shoyer @jekbradbury 

---

Updated: thanks to @jekbradbury, I am able to make this work. This PR resolves #1789, https://github.com/google/jax/issues/1875#issuecomment-567026911, #1188, #552 (actually, the fourth one is not relevant, but with the change in this PR, it seems fast to get samples from gamma sampler now).

@mattjj It would be great to have this for the next jax release. :)

Problem statement:
grad and vmap do not composable with gamma sampler
Here is a repro code, which triggers error `NotImplementedError: Forward-mode differentiation rule for 'while' not implemented` when trying to take gradient of `g`.
```
import jax
import jax.numpy as np
jax.config.update('jax_platform_name', 'cpu')

def f(a):
    return jax.random.gamma(jax.random.PRNGKey(0), a)

def g(x):
    return np.sum(jax.vmap(f)(x))

print(g(np.ones(3)))
print(jax.grad(g)(np.ones(3)))
```

Hints from issue discussion:
This is another manifestation of https://github.com/google/jax/issues/1249.
Oh... Thanks, @shoyer! That seems like the problem which I get. I'll try to see if I can go back to several months ago to not use custom_transform for gamma sampler.

Commit message: not use custom transform for gamma sampler

Patch to review:
diff --git a/jax/random.py b/jax/random.py
index 92806f1c7308..941a0ddec2e8 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -41,7 +41,9 @@
 from jax import abstract_arrays
 from jax.scipy.special import logit
 from jax.scipy.linalg import cholesky
+from jax.interpreters import ad
 from jax.interpreters import batching
+from jax.interpreters import partial_eval as pe
 from jax.interpreters import xla
 
 
@@ -878,18 +880,46 @@ def _case4(zagf):
 def _gamma_grad(sample, a):
     samples = np.reshape(sample, -1)
     alphas = np.reshape(a, -1)
-    grads = vmap(_gamma_grad_one)(samples, alphas)
+    if np.size(alphas) == 1:
+        grads = _gamma_grad_one(samples[0], alphas[0])
+    else:
+        # TODO: benchmark execute time against grads = vmap(_gamma_grad_one)(samples, alphas)
+        grads = lax.map(lambda args: _gamma_grad_one(*args), (samples, alphas))
     return grads.reshape(onp.shape(a))
 
-@custom_transforms
 def _gamma_impl(key, a):
-    alphas = np.reshape(a, -1)
-    keys = split(key, onp.size(alphas))
-    samples = vmap(_gamma_one)(keys, alphas)
-    return np.reshape(samples, onp.shape(a))
-
-defjvp(_gamma_impl, None,
-       lambda tangent, ans, key, a, **kwargs: tangent * _gamma_grad(ans, a))
+  if key.ndim == 2:  # batch of keys and alphas
+    size = np.size(a[0])
+    if size > 1:
+      key = lax.map(lambda k: split(k, size), key)
+  else:
+    size = np.size(a)
+    if size > 1:
+      key = split(key, size)
+  alphas = np.reshape(a, -1)
+  keys = np.reshape(key, (-1, 2))
+  if np.size(alphas) == 1:
+    samples = _gamma_one(keys[0], alphas[0])
+  else:
+    # XXX in GPU, using lax.map is slower than using vmap if alphas.size > 50000
+    # but that usage case is rare and can be resolved by vectorizing gamma sampler
+    samples = lax.map(lambda args: _gamma_one(*args), (keys, alphas))
+  return np.reshape(samples, np.shape(a))
+
+def _gamma_batching_rule(batched_args, batch_dims):
+    k, a = batched_args
+    bk, ba = batch_dims
+    size = next(t.shape[i] for t, i in zip(batched_args, batch_dims) if i is not None)
+    k = batching.bdim_at_front(k, bk, size)
+    a = batching.bdim_at_front(a, ba, size)
+    return random_gamma_p.bind(k, a), 0
+
+random_gamma_p = core.Primitive('random_gamma')
+random_gamma_p.def_impl(_gamma_impl)  # partial(xla.apply_primitive, random_gamma_p))
+random_gamma_p.def_abstract_eval(lambda key, a: abstract_arrays.raise_to_shaped(a))
+ad.defjvp2(random_gamma_p, None, lambda tangent, ans, key, a, **kwargs: tangent * _gamma_grad(ans, a))
+xla.translations[random_gamma_p] = xla.lower_fun(_gamma_impl, instantiate=True)
+batching.primitive_batchers[random_gamma_p] = _gamma_batching_rule
 
 def gamma(key, a, shape=None, dtype=onp.float64):
   """Sample Gamma random values with given shape and float dtype.
@@ -911,7 +941,7 @@ def gamma(key, a, shape=None, dtype=onp.float64):
   dtype = dtypes.canonicalize_dtype(dtype)
   return _gamma(key, a, shape, dtype)
 
-@partial(jit, static_argnums=(2, 3))
+# @partial(jit, static_argnums=(2, 3))
 def _gamma(key, a, shape, dtype):
   if shape is None:
     shape = onp.shape(a)
@@ -921,7 +951,7 @@ def _gamma(key, a, shape, dtype):
   a = lax.convert_element_type(a, dtype)
   if onp.shape(a) != shape:
     a = np.broadcast_to(a, shape)
-  return _gamma_impl(key, a)
+  return random_gamma_p.bind(key, a)
 
 
 def gumbel(key, shape=(), dtype=onp.float64):
diff --git a/tests/random_test.py b/tests/random_test.py
index f025626e5a32..3ace0c05c3fd 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -27,10 +27,12 @@
 import scipy.stats
 
 from jax import api
+from jax import grad
 from jax import lax
 from jax import numpy as np
 from jax import random
 from jax import test_util as jtu
+from jax import vmap
 from jax.interpreters import xla
 
 from jax.config import config
@@ -443,6 +445,12 @@ def testIssue756(self):
     else:
       self.assertEqual(onp.result_type(w), onp.float32)
 
+  def testIssue1789(self):
+    def f(x):
+      return random.gamma(random.PRNGKey(0), x)
+
+    grad(lambda x: np.sum(vmap(f)(x)))(np.ones(2))
+
   def testNoOpByOpUnderHash(self):
     def fail(*args, **kwargs): assert False
     apply_primitive, xla.apply_primitive = xla.apply_primitive, fail


Code review comments:
diff_hunk: @@ -878,18 +880,46 @@ def _case4(zagf):
 def _gamma_grad(sample, a):
     samples = np.reshape(sample, -1)
     alphas = np.reshape(a, -1)
-    grads = vmap(_gamma_grad_one)(samples, alphas)
+    if np.size(alphas) == 1:
+        grads = _gamma_grad_one(samples[0], alphas[0])
+    else:
+        # TODO: benchmark execute time against grads = vmap(_gamma_grad_one)(samples, alphas)
+        grads = lax.map(lambda args: _gamma_grad_one(*args), (samples, alphas))
     return grads.reshape(onp.shape(a))
 
-@custom_transforms
 def _gamma_impl(key, a):
-    alphas = np.reshape(a, -1)
-    keys = split(key, onp.size(alphas))
-    samples = vmap(_gamma_one)(keys, alphas)
-    return np.reshape(samples, onp.shape(a))
-
-defjvp(_gamma_impl, None,
-       lambda tangent, ans, key, a, **kwargs: tangent * _gamma_grad(ans, a))
+  if key.ndim == 2:  # batch of keys and alphas
+    size = np.size(a[0])
+    if size > 1:
+      key = lax.map(lambda k: split(k, size), key)
+  else:
+    size = np.size(a)
+    if size > 1:
+      key = split(key, size)
+  alphas = np.reshape(a, -1)
+  keys = np.reshape(key, (-1, 2))
+  if np.size(alphas) == 1:
+    samples = _gamma_one(keys[0], alphas[0])
+  else:
+    # XXX in GPU, using lax.map is slower than using vmap if alphas.size > 50000
+    # but that usage case is rare and can be resolved by vectorizing gamma sampler
+    samples = lax.map(lambda args: _gamma_one(*args), (keys, alphas))
+  return np.reshape(samples, np.shape(a))
+
+def _gamma_batching_rule(batched_args, batch_dims):
+    k, a = batched_args
+    bk, ba = batch_dims
+    size = next(t.shape[i] for t, i in zip(batched_args, batch_dims) if i is not None)
+    k = batching.bdim_at_front(k, bk, size)
+    a = batching.bdim_at_front(a, ba, size)
+    return random_gamma_p.bind(k, a), 0
+
+random_gamma_p = core.Primitive('random_gamma')
+random_gamma_p.def_impl(_gamma_impl)  # partial(xla.apply_primitive, random_gamma_p))
path: jax/random.py
text: @author:
I believe I should use `partial(xla.apply_primitive, random_gamma_p))` here but it doesn't work.