Whats the best way to map across a factorial function?

[Peter-Vincent]

I want to construct a factorial function. Doing it recursively is the standard method, but doesn't play well when I try and map across it, since we need to use the value for control flow (see example code below)

def _factorial(self,value):
        if value ==0:
            return value
        return value * self._factorial(value - 1)

self.factorial = vmap(self._factorial)

Instead I tried the following code

def _factorial(self,value):
        return lax.reduce(jnp.arange(1,value+1,dtype=float),1,jnp.multiply,[0])
self.factorial = vmap(self._factorial)

But this still yields a very similar error

return lax.reduce(jnp.arange(1,value+1,dtype=float),1,jnp.multiply,0)
  File "/home/peter/anaconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3048, in arange
    stop = None if stop is None else require(stop, msg("stop"))
jax._src.traceback_util.FilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([ 1,  1,  1,  1,  1,  1,  1,  1,  2,  7, ...
                                                                                       1,  1], dtype=int32)
       batch_dim = 0
It arose in jax.numpy.arange argument `stop`.

I've only included some of the elements of the DeviceArray here. Then of course there are the internal frames.

I guess form this jnp.arange uses a boolean to decide when to stop iterating.

If I instead use jnp.linspace(0,value-1,value)+1 or jnp.array(range(0,value-1)+1 I get the same output, which makes me think that Python is, under the hood, turning the value argument into something that Jax doesn't play well with.

Any ideas on how to get around this?

[Peter-Vincent]

The following code works

def _factorial(self,value):
        init_val = (value,1)
        cond_fun = lambda x: x[0] > 1
        body_fun = lambda x: (x[0]-1,x[1] * x[0])
        return lax.while_loop(cond_fun,body_fun,init_val)[1]

But I'm still wondering whats going under the hood in jnp.arange and jnp.linspace that means they don't play nice.

Thanks!

[jakevdp]

The reason you're running into issues with arange and linspace is that you are constructing arrays with data-dependent shapes. In JIT-compiled or vmapped code, all arrays must have static shapes.

If you want a fast, vmappable factorial function, the best approach is probably to use jax.scipy.special.gammaln: this computes the log of the gamma function, which is related to the factorial function. For example:

import jax.numpy as jnp
from jax.scipy.special import gammaln

def factorial(n):
  return jnp.exp(gammaln(n + 1))

print(factorial(jnp.arange(10)))
# [1.0000005e+00 1.0000000e+00 1.9999986e+00 6.0000024e+00 2.4000023e+01
#  1.2000006e+02 7.2000177e+02 5.0400068e+03 4.0319996e+04 3.6288172e+05]

[Peter-Vincent]

that makes sense. Thank you!