Repository: google/jax

Issue:
TypeError: Cannot determine dtype of Zero(ShapedArray(complex128[]))
### Description

Hi,

My code has complex numbers and backprop.

I used `jax.value_and_grad()` function to get gradient but got TypeError
`TypeError: Cannot determine dtype of Zero(ShapedArray(complex128[121]))`

It came from `.conj()` function so I replaced it with manual conjugation and it worked. 

So in summary,

1. jax.value_and_grad
    1. jit, with `arr.conj()`: Type Error
    2. jit, with `arr.real + arr.imag * -1j`: OK
    3. no jit, with `arr.conj()` : OK
    4. no jit, with `arr.real + arr.imag * -1j`: OK

Can you please take a look?

> /Users/yongha/miniconda3/envs/meent/bin/python /Users/yongha/project/rcwa/meent/QA/grad_complex.py 
> Traceback (most recent call last):
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/dtypes.py", line 611, in dtype
>     dt = np.result_type(x)
>   File "<__array_function__ internals>", line 200, in result_type
> TypeError: Cannot interpret 'Zero(ShapedArray(complex128[121]))' as a data type
> 
> The above exception was the direct cause of the following exception:
> 
> Traceback (most recent call last):
>   File "/Users/yongha/project/rcwa/meent/QA/grad_complex.py", line 66, in <module>
>     grad = jmee.grad(pois, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 53, in grad
>     _, grads = self._grad(params, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 48, in _grad
>     loss_value, grads = jax.value_and_grad(forward_pass)(params, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 44, in forward_pass
>     result = forward(**params)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/_base.py", line 217, in wrap
>     res = func(*args, **kwargs)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/rcwa.py", line 196, in conv_solve
>     de_ri, de_ti, layer_info_list, T1, kx_vector = self._conv_solve_jit()
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/rcwa.py", line 150, in _conv_solve_jit
>     return self._conv_solve()
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/rcwa.py", line 145, in _conv_solve
>     de_ri, de_ti, layer_info_list, T1, kx_vector = self._solve(self.wavelength, E_conv_all, o_E_conv_all)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/rcwa.py", line 114, in _solve
>     de_ri, de_ti, layer_info_list, T1 = self.solve_2d(wavelength, e_conv_all, o_e_conv_all)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/_base.py", line 217, in wrap
>     res = func(*args, **kwargs)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/_base.py", line 389, in solve_2d
>     = transfer_2d_1(ff_x, ff_y, ff_xy, k0, self.n_I, self.n_II, self.kx_vector, self.period, fourier_indices_y,
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/transfer_method.py", line 272, in transfer_2d_1
>     k_I_z = conj(k_I_z)  # manual conjugate
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/emsolver/primitives.py", line 9, in conj
>     return arr.conj()
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/numpy/ufuncs.py", line 599, in conjugate
>     return lax.conj(x) if np.iscomplexobj(x) else _asarray(x)
> jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: Cannot determine dtype of Zero(ShapedArray(complex128[121]))
> 
> The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
> 
> --------------------
> 
> The above exception was the direct cause of the following exception:
> 
> Traceback (most recent call last):
>   File "/Users/yongha/project/rcwa/meent/QA/grad_complex.py", line 66, in <module>
>     grad = jmee.grad(pois, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 53, in grad
>     _, grads = self._grad(params, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 48, in _grad
>     loss_value, grads = jax.value_and_grad(forward_pass)(params, forward, loss_fn)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
>     return fun(*args, **kwargs)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/api.py", line 741, in value_and_grad_f
>     g = vjp_py(lax_internal._one(ans))
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/tree_util.py", line 303, in __call__
>     return self.fun(*args, **kw)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/api.py", line 2183, in _vjp_pullback_wrapper
>     ans = fun(*args)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/tree_util.py", line 303, in __call__
>     return self.fun(*args, **kw)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 146, in unbound_vjp
>     arg_cts = backward_pass(jaxpr, reduce_axes, True, consts, dummy_args, cts)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 250, in backward_pass
>     cts_out = reducing_transposes[eqn.primitive](
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/pjit.py", line 1740, in _pjit_transpose
>     transpose_jaxpr = _pjit_transpose_trace(body, global_cts_in_avals)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/linear_util.py", line 322, in memoized_fun
>     ans = call(fun, *args)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/pjit.py", line 1717, in _pjit_transpose_trace
>     transpose_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(fun, in_avals)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
>     return func(*args, **kwargs)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2049, in trace_to_jaxpr_dynamic
>     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2066, in trace_to_subjaxpr_dynamic
>     ans = fun.call_wrapped(*in_tracers_)
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/linear_util.py", line 166, in call_wrapped
>     ans = self.f(*args, **dict(self.params, **kwargs))
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 264, in closed_backward_pass
>     return backward_pass(jaxpr.jaxpr, reduce_axes, transform_stack, jaxpr.consts,
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 253, in backward_pass
>     cts_out = get_primitive_transpose(eqn.primitive)(
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1895, in _conj_transpose_rule
>     return [conj(t)]
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 350, in conj
>     return conj_p.bind(x, input_dtype=_dtype(x))
>   File "/Users/yongha/miniconda3/envs/meent/lib/python3.10/site-packages/jax/_src/dtypes.py", line 613, in dtype
>     raise TypeError(f"Cannot determine dtype of {x}") from err
> jax._src.traceback_util.UnfilteredStackTrace: TypeError: Cannot determine dtype of Zero(ShapedArray(complex128[121]))
> 
> The stack trace below excludes JAX-internal frames.
> The preceding is the original exception that occurred, unmodified.
> 
> --------------------
> 
> The above exception was the direct cause of the following exception:
> 
> Traceback (most recent call last):
>   File "/Users/yongha/project/rcwa/meent/QA/grad_complex.py", line 66, in <module>
>     grad = jmee.grad(pois, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 53, in grad
>     _, grads = self._grad(params, forward, loss_fn)
>   File "/Users/yongha/project/rcwa/meent/meent/on_jax/optimizer/optimizer.py", line 48, in _grad
>     loss_value, grads = jax.value_and_grad(forward_pass)(params, forward, loss_fn)
> TypeError: Cannot determine dtype of Zero(ShapedArray(complex128[121]))
> 
> Process finished with exit code 1
> 

### What jax/jaxlib version are you using?

jax 0.4.1,  jaxlib 0.4.1

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10, M1 Mac

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! Could you provide a runnable repro?
Actually, I think I see the issue.
That's great. I tried to make a runnable code but wasn't that easy.
So if you can clone my repo and run one file, it will be faster. It that's not an option, I'll try to make sample code.
Please let me know.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -355,6 +355,7 @@ def complex(x: ArrayLike, y: ArrayLike) -> Array:
 
 def conj(x: ArrayLike) -> Array:
   r"""Elementwise complex conjugate function: :math:`\overline{x}`."""
+  # TODO(mattjj): remove input_dtype, not needed anymore
   return conj_p.bind(x, input_dtype=_dtype(x))
 
 def abs(x: ArrayLike) -> Array:
@@ -1896,7 +1897,9 @@ def _conj_impl(x, **kw):
 
 def _conj_transpose_rule(t, x, *, input_dtype):
   assert ad.is_undefined_primal(x)
-  if dtypes.issubdtype(input_dtype, np.complexfloating):
+  if type(t) is ad_util.Zero:
+    return [ad_util.Zero(x.aval)]
+  elif dtypes.issubdtype(input_dtype, np.complexfloating):
     return [conj(t)]
   else:
     return [real(t)]


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4214,6 +4214,12 @@ def outer_fn(x):
     self.assertEqual(inner_count, 1)
     self.assertEqual(outer_count, 1)
 
+  def test_grad_conj_symbolic_zeros(self):
+    # https://github.com/google/jax/issues/15400
+    f = lambda x: jax.jit(lambda x, y: (x, y))(x, jax.lax.conj(x))[0]
+    out = jax.grad(f)(3.0)  # doesn't crash
+    self.assertAllClose(out, 1., check_dtypes=False)
+
 
 class RematTest(jtu.JaxTestCase):