Why jvp imposes primals and tangents to be on the same dtype?

[mariogeiger]

Why is it not possible to compute the following?

import jax
import jax.numpy as jnp

jax.jvp(jnp.sin, (0.1,), (1. + 1j,))  # ERROR

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 float32 and so expected tangent dtype float32, but got tangent dtype complex64 instead.

[jakevdp]

In general JAX does not do implicit type promotion in binary operations (e.g. try running jax.lax.add(1, 1.0)) and this convention extends to JVP rules. This kind of stricter convention can help catch bugs, and so it's useful to enforce for lower-level APIs.

The jax.numpy package, of course, does do automatic type promotion for many operations (e.g. jnp.add(1, 1.0)), mainly because numpy itself supports this kind of operation and jax.numpy mirrors that API.