Pure Callbacks do not support JVP. Is there a workaround?

[kgierach-hf]

Hi,
I'm new to JAX and using it within the context of numpyro. I would like to perform a data transformation based on weights, and a shift quantity computed from a sampled variable. I was unable to get the code to run since the computed statistic became a Tracer, and therefore is unable to direct control flow. So, I've gone the route of attempting to implement the transformation/convolution using callbacks.

the error message is:
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Pure callbacks do not support JVP. Please use jax.custom_jvp to use callbacks while taking gradients.

Is there a workaround to avoid implementing the custom_jvp function?

Next, I have attempted a custom_jvp function as shown in the documentation, but that function is presumably not being called, and the same error message is emitted as above. The steps I followed are:

1. added the @custom_jvp annotation to the original callback function named "compute_weights_host".
2. added the @compute_weights_host.defjvp annotation to the JVP function.

Am I missing anything else?

[jakevdp]

The workaround is to define the custom_jvp function. In some cases, you may need to define custom_vjp instead (there's some discussion in the Custom derivative rules docuementation).

Can you say more about what error you're seeing when you try this? Can you share a reproduction of the issue?

[kgierach-hf]

Yes! I'm working on a simple example to reproduce what I'm experiencing. Thank you @jakevdp

[kgierach-hf]

Here is an example that reproduces the problem @jakevdp . Happy to send you the notebook directly if that helps.

import jax
import jax.numpy as jnp
from jax import random, vmap
from jax import custom_jvp
import numpyro
import numpy as np
import math
import pandas as pd

import numpyro
from numpyro import handlers
from numpyro.diagnostics import hpdi
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
import functools

### reverse adstock example attempting to leverage JAX callbacks
### data gen: ground truth + transformed data with shift and decay
shift = 2.4
decay = 0.65
weights = np.power( decay, np.arange(13) )   # most decay occurs within 13 weeks
weights = weights / np.sum( weights )

np.random.seed(101)
print( "shift floor: ", shift-math.floor(shift) , "|", (1.0 - (shift - math.floor(shift)) ) )

weights_shifted = np.zeros(13)
for wIdx in range(0,len(weights)):
    if math.floor(shift)+wIdx+1 < len(weights):
        weights_shifted[ math.floor(shift)+wIdx   ] += weights[ wIdx ] * (1.0 - (shift - math.floor(shift)))
        weights_shifted[ math.floor(shift)+wIdx+1 ] += weights[ wIdx ] * (shift-math.floor(shift))

npd = np.array( np.random.normal( loc=100.0, scale=25, size=100 ) )
npd_xformed = np.zeros( 100 + 13 + 3 )
for idx in range(0,len(npd)):
    val = npd[idx]
    vec = np.ones( len(weights_shifted) ) * val
    left_pad  = idx
    right_pad = len(npd_xformed) - left_pad - len(weights_shifted)
    mul_weights = np.multiply( vec, weights_shifted )
    padded = np.pad( mul_weights, (left_pad,right_pad) )
    npd_xformed = np.add( npd_xformed, padded ) 
    observed = npd_xformed[0:len(npd)]

# train the model , reproduce issue
df = pd.DataFrame()
df['ground_truth'] = npd
df['obs']          = observed 

### performs the work of distributing the spend according to the weights and the provided shift value.


@custom_jvp
def distribute_weights( shifted_weights, mean_week, colData, colDataLen, weeks ):
    print( 'distribute_weights(): BEGIN' )

    bVerbose = False

    rev_shift_start = math.ceil( mean_week )  # week index where spend is carried forward
    colDataOut = np.zeros( np.shape(colData), dtype=np.float32 )

    for dIdx in range( 0, colDataLen ):
        mul_data    = np.ones( weeks, dtype=np.float32 )
        spnd_data   = mul_data * colData[dIdx]
        split_spend = np.multiply( spnd_data, shifted_weights )

        # find start index of xform_data
        sIdx = max( (-1*rev_shift_start) + dIdx, 0 )
        # find start index of weights
        weightsIdx = np.where( dIdx >= rev_shift_start, 0, (rev_shift_start - dIdx) )

        eIdx = (sIdx + weeks - weightsIdx)
        if bVerbose:
            print( '    sIdx: ', sIdx, ' -> eIdx: ', eIdx )
        weightsEnd = -1

        if eIdx > colDataLen:
            weightsEnd = len(split_spend) - (eIdx-colDataLen)

        split_spend  = split_spend[weightsIdx:weightsEnd]
        leftPad      = sIdx
        rightPad     = colDataLen - sIdx - len(split_spend)
        padded_spend = np.pad( split_spend, (leftPad,rightPad) )
        colDataOut   = np.add( colDataOut, padded_spend )

    return np.array( colDataOut, dtype=np.float32 )

### define custom JVP for this callback 
dw = jax.custom_jvp( distribute_weights )


### JVP callback
@distribute_weights.defjvp
def distribute_weights_host_jvp( primals, tangents ):
    print( "distribute_weights_host_jvp(): primals: ", primals )
    print( "distribute_weights_host_jvp(): tangents: ", tangents )
    shifted_weights, mean_week, colData, colDataLen, weeks = primals
    x_dot = tangents
    primal_out = distribute_weights( shifted_weights, mean_week, colData, colDataLen, weeks )
    tangent_out = -1 * jnp.multiply( primal_out, x_dot )
    return primal_out, tangent_out



### implemented shift logic as JAX Callback since the shift amount comes out of a sampled variable.
def distribute_weights_wrapper( shifted_weights, mean_week, colData, colDataLen, weeks ):
    result_shape = jax.ShapeDtypeStruct( jnp.shape(colData), colData.dtype )
    return jax.pure_callback( dw, result_shape, shifted_weights, mean_week, colData, colDataLen, weeks )


### current problem is with this function
@custom_jvp
def compute_weights_host( lam ):

    pct=0.5
    scaler = 1.0
    max_week=13.0
    max_horizon=100

    bVerbose = False

    if bVerbose:
        print( 'CWH(): BEGIN METHOD' )

    def geo_f_local( lam, t, scaler = 1.0 ):
        return np.exp( -1.0 * scaler * lam * t )

    ret_val = np.zeros( int(max_week+1), dtype=np.float32 )
    norm_const = integrate.quad( geo_f_local, 0, max_horizon, args=(lam) )[0]
    # print( "norm_const: ", norm_const )
    t = np.arange( 0.0, max_week, 0.1 )
    res = max_week
    for x in range(len(t)):
        v = (integrate.quad( geo_f_local, 0, t[x], args=(lam, scaler) )[0] / norm_const)
        if v > (1.0-pct):
            res = t[x]
            break

    mean_point = res
    mean_time = res
    cur_week = math.ceil( res )
    for weekIdx in range( cur_week, 0, -1 ):
        startPt = max( res - 1.0, 0 )
        # print( "weekIdx[", weekIdx, "]: ", startPt,  " | ", res )
        v = (integrate.quad( geo_f_local, startPt, res, args=(lam, scaler) )[0] / norm_const)
        res = res - 1.0
        ret_val[ weekIdx-1 ] = v

    for weekIdx in range( cur_week, int(max_week) ):
        endPt = (mean_time+1) if weekIdx < (max_week-1) else max_horizon
        v = (integrate.quad( geo_f_local, mean_time, endPt, args=(lam, scaler) )[0] / norm_const)
        mean_time = mean_time + 1.0
        ret_val[ weekIdx ] = v

    ret_val[ int(max_week) ] = mean_point
    if bVerbose:
        print( 'CWH(): ', ret_val )
    return ret_val


### define custom JVP
cwh = jax.custom_jvp( compute_weights_host )

### JVP callback for computing weights
@compute_weights_host.defjvp
def compute_weights_host_jvp( primals, tangents ):
    print( "compute_weights_host_jvp(): primals: ", primals )
    print( "compute_weights_host_jvp(): tangents: ", tangents )
    lam   = primals
    x_dot = tangents
    primal_out = compute_weights_host( lam )
    tangent_out = -1 * lam * jnp.multiply( primal_out, x_dot )
    return primal_out, tangent_out


### wrap callback
def compute_weights_wrapper( lam ):
    result_shape = jax.ShapeDtypeStruct( (13+1,), lam.dtype )
    return jax.pure_callback( cwh, result_shape, lam )


@functools.partial( jax.jit, static_argnames=["num_weeks", "colDataLen"])
def run_reverse_shift( lag_weight, colData, num_weeks, colDataLen ):
    print(  "run_reverse_shift: lag_weight: ", lag_weight, ":", type(lag_weight) )
    shifted_weights = compute_weights_wrapper( lag_weight )
    
    mean_week = shifted_weights[-1]
    print( 'mean week x: ', mean_week )
    weights_only = shifted_weights[0:13]
    print( 'weights only: ', shifted_weights, "|", jnp.shape(shifted_weights) )

    colData = distribute_weights_wrapper( weights_only, mean_week, colData, colDataLen, num_weeks )
    # print( "run_reverse_shift: ", colData, " | ", jnp.shape( colData ) )
    return colData


def transform( data: jnp.ndarray ):
    alpha = numpyro.sample( "alpha",  dist.Beta(2,2) )
    ret = run_reverse_shift( alpha, data, 13, len(data) )
    return ret
    

def model( measurements=None, ground_truth=None ):
    jnp_measurements = jnp.array(measurements)
    jnp_ground_truth = jnp.array(ground_truth)
    mu = numpyro.deterministic( name="mu", value=transform( jnp_measurements ) )
    sigma = numpyro.sample("sigma", dist.Exponential(1.0))
    numpyro.sample( "obs", dist.Normal(mu, sigma), obs=jnp_ground_truth )
    return


### Start from this source of randomness. We will split keys for subsequent operations.
rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)

### Run NUTS.
kernel = NUTS(model)
num_samples = 2000
mcmc = MCMC( kernel, num_warmup=1000, num_samples=num_samples )
mcmc.run(
    rng_key_, measurements=df['obs'].values, ground_truth=df['ground_truth'].values
)
mcmc.print_summary()
samples_1 = mcmc.get_samples()

Error message:

File "/Users/karl.gierach/Python/py310-lightweightmmm/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 184, in reraise_with_filtered_traceback
  File "/Users/karl.gierach/Python/py310-lightweightmmm/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 242, in __call__
AttributeError: No JVP defined for custom_jvp function compute_weights_host using defjvp.

[jakevdp]

There are two problems here I think:

(1) you define cwh as a custom_jvp, but you never define its JVP via cwh.defjvp. Fortunately you define its JVP elsewhere, so you can delete cwh and just use compute_weights_host directly.
(2) you use compute_weights_wrapper in a setting where it is automatically differentiated, but it calls jax.pure_callback which is not differentiable. You'll need to define a custom_jvp rule for compute_weights_wrapper to avoid trying to automatically differentiate the pure_callback.

You might find this doc helpful: it contains a worked example of a pure callback using custom_jvp to define autodiff rules: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html#example-pure-callback-with-custom-jvp Note that the custom_jvp is defined for the function that calls the pure_callback, not the function passed to the pure callback.

Hope that helps!

[kgierach-hf]

Thanks @jakevdp ! Your tip helped me get the code up and running. I restructured the code as you described above, and provided the custom_jvp implementation to the function which invokes the pure callback.