How to solve the TypeError("primal and tangent arguments to jax.jvp do not match) when one of the primal's dtype is int32?

[AndrewLiu0725]

I have a neural network that takes discretized inputs, specifically an array of integers. When I attempted to use custom_jvp to craft a customized gradient (with respect to model parameters) calculation for my loss function, I encountered the error message:

TypeError: primal and tangent arguments to jax.jvp do not match; dtypes must be equal, or in case of int/bool primal dtype the tangent dtype must be float0.Got primal dtype int32 and so expected tangent dtype [('float0', 'V')], but got tangent dtype int32 instead. 

This error emerged because one of the primal's dtypes is int32. I am unsure how to resolve this type of mismatch error. Should I manually generate a float0 array that matches the shape of the primal/tangent and then use it in place of the jvp, given that I am only taking derivatives with respect to the model parameters? Any assistance is greatly appreciated.

Below is a minimal reproducible code that triggers the error:

import jax.numpy as jnp 
from jax import jit, vmap, value_and_grad, jvp, custom_jvp, random

def mini_NN(params, NN_input):
    return params[0]*NN_input + params[1]
mini_NN_v = vmap(mini_NN, in_axes=(None, 0))

key = random.PRNGKey(42)
init_params = (random.uniform(key, shape=(4, 1)), random.uniform(key, shape=(1,)))

init_NN_inputs = jnp.array([[1,-1,1,1], [-1,-1,1,-1]], dtype=jnp.int32)

@custom_jvp
@jit
def loss(params, NN_inputs):
    return jnp.mean(mini_NN_v(params, NN_inputs))

@loss.defjvp
@jit
def loss_jvp(primals, tangents):
    psi, psi_tangent = jvp(mini_NN_v, primals, tangents)
    primal_out = jnp.mean(psi)
    tangent_out = jnp.mean(psi_tangent**2) # my function was more complicated. Here I just use some random simple function
    return primal_out, tangent_out

print(loss(init_params, init_NN_inputs))
print(value_and_grad(loss, argnums=0)(init_params, init_NN_inputs)) # TypeError: primal and tangent arguments to jax.jvp do not match; dtypes must be equal, or in case of int/bool primal dtype the tangent dtype must be float0.Got primal dtype int32 and so expected tangent dtype [('float0', 'V')], but got tangent dtype int32 instead.

[jakevdp]

This looks similar to the issue addressed in #14570, but I'm not sure of the easiest way to do what you have in mind. I think passing symbolic_zeros=True to defjvp gets you part of the way there, but you can't forwarding this directly to the jvp call, and SymbolicZeros is not compatible with jit. I think @mattjj might have ideas about how to approach this.