complex -> real casting warning only raised during gradient computation of function using tensordot

[DanPuzzuoli]

I'm running into an issue where a complex -> real casting warning is being raised when computing the gradient of a function, but not when calling the function itself. The function I found this in involves computing tensordot on a pair of arrays with one being real and one complex. Below is a minimal example:

from jax import grad
import jax.numpy as jnp

def g(x):
    td = jnp.tensordot(x * jnp.array([1.]), jnp.array([3. + 1j], dtype=complex), axes=1)
    return td.real

grad(g)(1.)

grad(g)(1.) works and returns the correct result, however it raises the warning:

/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/lax/lax.py:2638: ComplexWarning: Casting complex values to real discards the imaginary part
  x_bar = _convert_element_type(x_bar, x.aval.dtype, x.aval.weak_type)

This warning however is not raised when simply calling g, e.g.:

g(1.)

In my full case I know one of the arrays is real ahead of time, but the other is user-specified and could be either real or complex. (Of course, in the minimal example, we could have cast the second array as real as the imaginary part is irrelevant to the function output, but in my full code this is only a single piece of a larger, potentially complex-typed computation.)

It seems like this warning should not be raised during gradient computation, but I'm not sure of the latest policy on real/complex type mixing in functions like tensordot.

[DanPuzzuoli]

Just wanted to draw attention to this again - anyone have any thoughts this?