Error when trying to implement custom derivative with Bessel function of the first kind

[alexandermm]

Hi, I am trying to implement a function that uses jv from the scipy library written with a wrapper to work in jax. When I run the code,
the error given is AttributeError: No JVP defined for custom_jvp function modJv using defjvp. But I did define it using @custom_jvp. What is the issue here? The code is shown below.

Thank you,
Alex

import jax.numpy as jnp
from jax import custom_jvp
import scipy
from functools import partial


def jv(v, z):
  v, z = jnp.asarray(v), jnp.asarray(z)

  assert jnp.issubdtype(v.dtype, jnp.integer)

  z = z.astype(jnp.result_type(float, z.dtype))

  _scipy_jv = lambda v, z: scipy.special.jv(v, z).astype(z.dtype)

  result_shape_dtype = jax.ShapeDtypeStruct(shape=jnp.broadcast_shapes(v.shape, z.shape), dtype=z.dtype)

  return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vectorized=True)

@custom_jvp
def modJv(z, v, a):
    return jv(v, a*jnp.sqrt(z))

modJv = jax.custom_jvp(modJv)

@modJv.defjvp
def _modJv_jvp(primals, tangents):
    v,a, z     = primals
    _,_, z_dot = tangents
    
    Jv, Jvm2, Jvp2 = modJv(z, v,a), modJv(z, v-2,a), modJv(z, v+2,a)
    djv_dz = a*a/8.0*((Jv+Jvm2)/(v-1) - (Jv+Jvp2)/(v+1))
    
    return modJv(z, v, a), z_dot * djv_dz

modj = partial(modJv, 1, 2.5)
print(jax.grad(modj)(2.3))

[JiaYaobo]

Hi several comments!

• First

@custom_jvp
def modJv(z, v, a):
    return jv(v, a*jnp.sqrt(z))

modJv = jax.custom_jvp(modJv)

You have already decorated modJv with custom_jvp, so line modeJv=jax.custom_jvp(modJv) is redundant, delete it and the original error will be resolved.

• Second

The result will be nan since v=1 and v-1 is placed on the denominator

• Final

Recommend use kwargs wih partial like modj = partial(modJv, v=2, a=2.5), otherwise, with modj = partial(modJv, 1, 2.5), modj(2.3) doesn't work

Hope it helps

[alexandermm]

Hi Jia,

Thank you for the prompt reply! You are right about all the comments; I ended up modifying the jv function to be basically do what the modJv function is doing since I won't be using the jv declaration on its own as shown below.

Is there any difference with using scipy.special.jv(v, a*np.sqrt(z)).astype(z.dtype) with the np sqrt() or can I use jnp.sqrt()?

Thank you,
Alex

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

import numpy as np

import time

from functools import partial


@jax.custom_jvp
def modJv(v,a,z):
  v, a, z = jnp.asarray(v), jnp.asarray(a), jnp.asarray(z)

  # Require the order v to be integer type: this simplifies the JVP rule below.
  assert jnp.issubdtype(v.dtype, jnp.integer)

  # Promote the input to inexact (float/complex), jnp.result_type() accounts for enable_x64 flag.
  z = z.astype(jnp.result_type(float, z.dtype))
  a = a.astype(jnp.result_type(float, a.dtype))

  # Wrap scipy function to return the expected dtype.
  _scipy_jv = lambda v, a, z: scipy.special.jv(v, a*np.sqrt(z)).astype(z.dtype)

  # Define the expected shape & dtype of output.
  result_shape_dtype = jax.ShapeDtypeStruct(
      shape=jnp.broadcast_shapes(v.shape, a.shape, z.shape), dtype=z.dtype)

  # We use vectorize=True because scipy.special.jv handles broadcasted inputs.
  return jax.pure_callback(_scipy_jv, result_shape_dtype, v, a, z, vectorized=True)


@modJv.defjvp
def _modJv_jvp(primals, tangents):
  v, a, z     = primals
  _, _, z_dot = tangents
    
  Jv, Jvm2, Jvp2 = modJv(v,a, z), modJv(v-2,a, z), modJv(v+2,a, z)

  dmodJv_dz = a*a/8.0 * ((Jv+Jvm2)/(v-1) - (Jv+Jvp2)/(v+1))
    
  return modJv(v,a, z), z_dot * dmodJv_dz


v = 4
a = jnp.array(2.5, dtype=jnp.double)
z = jnp.array(1.0, dtype=jnp.double)
j1 = partial(modJv, v,a)
print(j1(z))
print(jax.grad(j1)(z))

[JiaYaobo]

Is there any difference with using scipy.special.jv(v, a*np.sqrt(z)).astype(z.dtype) with the np sqrt() or can I use jnp.sqrt()?

Sure, you can use jnp.sqrt, I believe the difference between numpy and jax.numpy has been comprehensively discussed in their document https://jax.readthedocs.io/en/latest/faq.html, but for this example, jax.pure_callback claims cleearly that pure_callback enables calling a Python function in JIT-ed JAX functions. The input callback will be passed NumPy arrays in place of JAX arrays and should also return NumPy arrays. Execution takes place on CPU, like any Python+NumPy function.

And just three different examples to explain the difference all running on cpu backend

• use np.sqrt in _scipy_jv

j1 = partial(jax.jit(modJv), v, a)
x = jnp.ones((1000,))
%timeit j1(x).block_until_ready()

117 µs ± 1.56 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

• use jnp.sqrt in _scipy_jv

j1 = partial(jax.jit(modJv), v, a)
x = jnp.ones((1000,))
%timeit j1(x).block_until_ready()

261 µs ± 6.33 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

• use jnp.sqrt out of _scipy_jv and modify codes

@jax.custom_jvp
def modJv(v,a,z):
  v, a, z = jnp.asarray(v), jnp.asarray(a), jnp.asarray(z)

  # Require the order v to be integer type: this simplifies the JVP rule below.
  assert jnp.issubdtype(v.dtype, jnp.integer)

  # Promote the input to inexact (float/complex), jnp.result_type() accounts for enable_x64 flag.
  z = z.astype(jnp.result_type(float, z.dtype))
  a = a.astype(jnp.result_type(float, a.dtype))

  # Wrap scipy function to return the expected dtype.
  a_sqrtz = a * jnp.sqrt(z)
  def _scipy_jv(v, a_sqrtz):
    return scipy.special.jv(v, a_sqrtz).astype(z.dtype)

  # Define the expected shape & dtype of output.
  result_shape_dtype = jax.ShapeDtypeStruct(
      shape=jnp.broadcast_shapes(v.shape, a.shape, z.shape), dtype=z.dtype)

  # We use vectorize=True because scipy.special.jv handles broadcasted inputs.
  return jax.pure_callback(_scipy_jv, result_shape_dtype, v, a_sqrtz, vectorized=True)

and then

j1 = partial(jax.jit(modJv), v, a)
x = jnp.ones((1000,))
%timeit j1(x).block_until_ready()

116 µs ± 1.31 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

I guess the slow performance of second condition is related to the cnversion from jnp.array to numpy.array in scipy.special.jv