float32 usage

[nsmith-]

What considerations prevent using float32? For now we have only this statement:

jaxfit/src/jaxfit/__init__.py

Lines 9 to 11 in e6b28d4

	# TODO: fine-grain control
	# Really the only place float32 falls apart is poisson logpmf for very large N
	_config.update("jax_enable_x64", True)

[nsmith-]

Trying to reproduce the claim in this note, I see what is probably a negligible numerical error:

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

def poisson_logpdf(x, mu):
    return special.xlogy(x, mu) - mu - special.gammaln(x+1)


@jax.custom_jvp
def poisson_logpdf_fancy(x, mu):
    """Standard log(P) for Poisson
    
    Implemented with custom jvp to improve accuracy
    when x and mu are both very large?
    """
    return special.xlogy(x, mu) - mu - special.gammaln(x+1)


@poisson_logpdf_fancy.defjvp
def _poisson_logpdf_jvp(primals, tangents):
    x, mu = primals
    dx, dmu = tangents
    df_dx = jnp.log(mu) - jax.grad(special.gammaln)(x+1)
    df_dmu = (x - mu)/mu
    return poisson_logpdf(x, mu), df_dx * dx + df_dmu * dmu


@jax.vmap
def compare(x, mu):
    return jax.grad(poisson_logpdf_fancy, 1)(x, mu) - jax.grad(poisson_logpdf, 1)(x, mu)


x = jnp.geomspace(1, 1e8, 20, dtype=jnp.float32)
mu = x + 1
fig, ax = plt.subplots()
ax.plot(compare(x, mu), label="Gradient diff")
ax.plot(x/mu - 1, label=r"$x/\mu-1$")
ax.set_ylim(-1e-7, 1e-7)
ax.legend()