Binomial coefficient in jax

[maurorigo]

Is there a direct way to compute a binomial coefficient which depends on some variable in a jax.lax.fori_loop (without using, for example, jax.numpy.exp() with jax.scipy.special.gammaln())? The standard scipy has scipy.special.comb(), though it can't be used as the variables in the fori_loop are jax tracer objects. Thank you in advance and sorry if the question is stupid.

[jakevdp]

I suspect the best way to compute binomial coefficients non-statically in JAX will be with gammaln; for example:

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

def comb(N, k):
  return jnp.exp(gammaln(N + 1) - gammaln(k + 1) - gammaln(N - k + 1))

print(jit(comb)(5, 3))
# 10.00001
print(scipy.special.comb(5, 3))
# 10

[maurorigo]

Yes that's the solution I went with, it's just slightly more cumbersome than a single function especially since I then need to cast the result into an int, but it's not a big deal. Thank you very much for your kind answer!