The dtype function does not work with Zero(ShapedArray(complex64[2,2]))

[rmoyard]

Hi Jax team! A question about jitting, I have a function f that returns a tuple of complex square Jax arrays and I want to get the jacobian of the imaginary part.

phi = jax.numpy.array(-0.792)

def f(phi):
    # get an initial sate
    init_state = ...

    # apply some unitary matrices on the state vector
    state =  ...

    density_matrix = jax.numpy.tensordot(state, jax.numpy.conj(state))


    # return a tuple of (densitty matrices, entropies, mutual information, ...)
    return tuple(measure(mp, state) for m in measurements)

def imag_out(phi):
    return tuple(jax.numpy.imag(r) for r in jax.jit(f)(phi))

grad_imag = jax.jacobian(imag_out)(phi)

I have a function that returns the image of a tuple of jax array, it works fine without jitting but when f is jitted I end up with this error. (important note: if I jit the jacobian instead it works as expected)

TypeError: Cannot interpret 'Zero(ShapedArray(complex64[2,2]))' as a data type

which happens in the dtype function

    def dtype(x: Any, *, canonicalize: bool = False) -> DType:
      """Return the dtype object for a value or type, optionally canonicalized based on X64 mode."""
      from jax._src import core     # TODO(frostig): break this cycle
      if x is None:
        raise ValueError(f"Invalid argument to dtype: {x}.")
      elif isinstance(x, type) and x in python_scalar_dtypes:
        dt = python_scalar_dtypes[x]
      elif type(x) in python_scalar_dtypes:
        dt = python_scalar_dtypes[type(x)]
      elif core.is_opaque_dtype(getattr(x, 'dtype', None)):
        dt = x.dtype
      else:
       dt = np.result_type(x)

It started failing with the version 0.4.4 and was working before.

EDIT

The issue arises when using jax.numpy.conj(state) which is 'Zero(ShapedArray(complex64[2,2]))'

    @_wraps(np.conjugate, module='numpy')
    @partial(jit, inline=True)
    def conjugate(x: ArrayLike, /) -> Array:
      _check_arraylike("conjugate", x)
>     return lax.conj(x) if np.iscomplexobj(x) else _asarray(x)
E     jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: Cannot interpret 'Zero(ShapedArray(complex64[2,2]))` as a data type
E     
E     The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
E     
E     --------------------

[jakevdp]

Thanks for the question! Can you edit your code snippet to define phi and f?

[rmoyard]

Hi @jakevdp , thanks a lot! Unfortunately the full example depends on another project, I am trying to come up with a stand alone example. But if it can help I have added some pseudo code.

Furthermore I am pretty sure the issue is arising when applying conj on the Zero(ShapedArray(complex64[2,2])) see the Edit