Anything faster than cond for splicing together functions?

[dieterichlawson]

I am writing a version of exprel = (exp(x) - 1)/x that has numerically stable forward and backward passes. When x=0, (exp(x)-1)/x is undefined because of division by zero. When x approaches 0 from either side, however, the function approaches 1. I would like to write a version that replaces the undefined portion with 1.

Using cond I can write:

def exprel(x):
  return jax.lax.cond(
      x == 0.,
      lambda _: 1.,
      lambda _x: jnp.expm1(_x) / _x,
      x)

My main question is: are there any better or faster ways to accomplish this? I will be calling this function millions of times, so I'd like to make sure it's fast.

[jakevdp]

This is fine as long as x is a scalar. If x is a vector you can use lax.select, or jnp.where which is a wrapper around that:

def exprel(x):
  return jnp.where(x == 0, 1, jnp.expm1(x) / x))

This is the pattern used throughout the JAX source code for this kind of thing.

Side-note, if you're interested in autodiff over this function, you may also need a custom_jvp to correctly define the derivative at x = 0; there's a tutorial on this here: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html

[jakevdp]

Regarding custom_jvp, this is similar enough to sinc that you might find that implementation interesting. To correctly compute Nth derivatives at x=0, it relies on a maclaurin series: https://github.com/google/jax/blob/961e09e61493842f5c057a2b2086813f2af58009/jax/_src/numpy/ufuncs.py#L714-L739

[dieterichlawson]

Awesome, thanks! If I vmap exprel, does that accomplish the same thing, or will where still be faster?

And yes, actually there are more serious issues with the gradient. My current approach is replacing the gradient in an epsilon ball around 0 with a taylor expansion of the gradient. I've included my code below in case it benefits others, with your suggested modifications.

def simple_exprel(x):
  return jnp.expm1(x) / x

grad_simple_exprel = jax.grad(simple_exprel)

@jax.custom_jvp
def exprel(x):
  return jnp.where(x == 0, 1, simple_exprel(x))

@exprel.defjvp
def exprel_jvp(xs, x_dots):
  x = xs[0]
  x_dot = x_dots[0]
  primal_out = exprel(x)
  tangent_out = jnp.where(
      jnp.less_equal(jnp.abs(x), 1e-2),
      0.5 + x/3. + jnp.square(x)/8.,
      grad_simple_exprel(x))
  return primal_out, tangent_out * x_dot

[dieterichlawson]

Just saw your response after posting mine, it looks like I arrived at a similar approach to sinc, thanks for the pointer!

[jakevdp]

vmap of cond lowers to select, so it's effectively the same thing.