More accurate experimental.jet

[YouJiacheng]

Since jit has constant folding, we can compute factorial with brute force method lax.prod(range(1, n + 1)) instead of approximate lax.exp(lax.lgamma(n+1.)).

Thanks to the great dynamic of python, we don't need to modify the source code of jax.

import jax
from jax.experimental import jet
jet.fact = lambda n: jax.lax.prod(range(1, n + 1)) # modification
def f(x: jnp.ndarray):
    return jnp.sum(jnp.square(x))

x = jnp.ones((3,), jnp.float32)
v = jnp.arange(3.)
print(jet.jet(f, (x,), ((v, jnp.zeros_like(x)),)))

With this modification, the result is
(DeviceArray(3., dtype=float32), [DeviceArray(6., dtype=float32), DeviceArray(10., dtype=float32)])
Without this modification, the result is
(DeviceArray(2.9999986, dtype=float32), [DeviceArray(5.999997, dtype=float32), DeviceArray(9.999995, dtype=float32)])
Look much better! (though maybe no difference for most application)

[mattjj]

That's a great idea. (Sorry for not noticing this until now.)

Want to make a PR to add this?

(Also, thanks for all the amazing contributions here recently!)

[YouJiacheng]

I don't know how to use brute force method only inside jit. Maybe this one-line "hack" is enough? Or add this hint into docstr of jet?