converting precision of JVP function produced by jax.linearize for mixed precision calculation

[ToshiyukiBandai]

Hi all,

I am interested in implementing mixed precision calculation for nonlinear solvers. In the following example, I want to calculate F in double precision and solve the Newton system in a single precision using Jacobian matrix computed in a single precision. Here, I have to compute F in double precision and JVP function in a single precision separately. Is there any way to get a value in double precision and JVP function in single precision? Is it possible to convert the double precision JVP function into single precision?

import jax
import jax.numpy as jnp
from jax import jvp, grad, jacrev

jax.config.update("jax_enable_x64", True)

def fun(x):
    return x**2 - jnp.asarray([1, 4])

x = jnp.asarray([1.1, 1.2])

F = fun(x) # double precision

_, f_jvp = jax.linearize(fun, jnp.float32(x)) # single precision

# the Jacobian matrix is efficiently constructed by lineax in practice
J_1 = f_jvp(jnp.asarray([1.0, 0.0], dtype = jnp.float32))
J_2 = f_jvp(jnp.asarray([0.0, 1.0], dtype = jnp.float32))
J = jnp.vstack((J_1, J_2)) # single precision

dx_mixed = jnp.linalg.solve(J, -jnp.float32(F))
dx_double = jnp.linalg.solve(jacrev(fun)(x), -F)

print(dx_mixed[0], dx_mixed[1])
print(dx_double[0], dx_double[1])