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_

Changed files:
jax/_src/lax/lax.py, tests/api_test.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 9c1edb8d1..d73ef2b12 100644
--- 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 @@ mlir.register_lowering(conj_p,
 
 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)]
diff --git a/tests/api_test.py b/tests/api_test.py
index ee736c4b1..06a6b28ba 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4214,6 +4214,12 @@ class APITest(jtu.JaxTestCase):
     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):