How to take complex conjugate of complex BCOO array

[DanPuzzuoli]

Hi,

As the title says, I'm just wondering what the "correct" way of taking the complex conjugate of a BCOO array is. Asking as sparsify(jnp.conjugate) does not seem to work.

[jakevdp]

Hi - there's not yet a sparse rule for conj (though it wouldn't be hard to add)

You could work around it by calling real and imag, which are defined:

sparse.sparsify(lambda x: jnp.real(x) - jnp.imag(x))(x)

I'll work on adding conj support in sparsify.

[jakevdp]

Support added in #18718

[DanPuzzuoli]

Thanks @jakevdp

[DanPuzzuoli]

Actually @jakevdp I just realized there is an issue with your suggested temporary solution:

y = BCOO.fromdense([1j, 0.])
sparsify(lambda x: jnp.real(x) - 1j * jnp.imag(x))(y)

seems to work, but:

y = BCOO.fromdense([1., 0.])
sparsify(lambda x: jnp.real(x) - 1j * jnp.imag(x))(y)

actually raises the error:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[36], line 2
      1 y = BCOO.fromdense([1., 0.])
----> 2 sparsify(lambda x: jnp.real(x) - 1j * jnp.imag(x))(y)

File /opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/experimental/sparse/transform.py:465, in _sparsify_with_interpreter.<locals>.wrapped(*args, **params)
    463 spenv = SparsifyEnv()
    464 spvalues = arrays_to_spvalues(spenv, args)
--> 465 spvalues_out, out_tree = f_raw(spenv, *spvalues, **params)
    466 out = spvalues_to_arrays(spenv, spvalues_out)
    467 return tree_unflatten(out_tree, out)

File /opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/experimental/sparse/transform.py:450, in sparsify_raw.<locals>.wrapped(spenv, *spvalues, **params)
    448 wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, params), in_tree)
    449 jaxpr, out_avals_flat, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals_flat)
--> 450 result = eval_sparse(jaxpr, consts, spvalues_flat, spenv)
    451 if len(out_avals_flat) != len(result):
    452   raise Exception("Internal: eval_sparse does not return expected number of arguments. "
    453                   "Got {result} for avals {out_avals_flat}")

File /opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/experimental/sparse/transform.py:427, in eval_sparse(jaxpr, consts, spvalues, spenv)
    425   if prim not in sparse_rules_bcoo:
    426     _raise_unimplemented_primitive(prim)
--> 427   out = sparse_rules_bcoo[prim](spenv, *invals, **eqn.params)
    428 else:
    429   out_bufs = prim.bind(*(spenv.data(val) for val in invals), **eqn.params)

File /opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/experimental/sparse/transform.py:652, in _sub_sparse(spenv, *spvalues)
    650   return _add_sparse(spenv, X, *sparse_rules_bcoo[lax.neg_p](spenv, Y))
    651 else:
--> 652   raise NotImplementedError("Subtraction between sparse and dense array.")

NotImplementedError: Subtraction between sparse and dense array.

So it seems that there is some issue if y has dtype float. Just fyi, seems relevant regardless of whether sparsify(jnp.conjugate) is implemented.

[jakevdp]

Yeah, the workaround assumes that the input is complex. If it's not, you should just return the input unmodified:

def conj_workaround(x):
  if jnp.issubdtype(x.dtype, jnp.complexfloating):
    return x.real - 1j * x.imag
  return x