Complex Vector Multiplication w/ Pallas

[bsaoptima]

from functools import partial

import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
import numpy as np


def cdot_vectors_kernel(x_real_ref, x_imag_ref, y_real_ref, y_imag_ref, o_ref):
    x_real = pl.load(x_real_ref, (slice(None),))
    x_imag = pl.load(x_imag_ref, (slice(None),))
    y_real = pl.load(y_real_ref, (slice(None),))
    y_imag = pl.load(y_imag_ref, (slice(None),))
    
    o_real = x_real * y_real - x_imag * y_imag
    o_imag = x_real * y_imag + x_imag * y_real

    o = jnp.array(o_real + 1j * o_imag)   
    pl.store(o_ref, (slice(None), ), o)


@jax.jit
def cdot_vectors(x_real: jax.Array, x_imag: jax.Array, y_real: jax.Array, y_imag: jax.Array) -> jax.Array:
    return pl.pallas_call(
        cdot_vectors_kernel,
        out_shape=jax.ShapeDtypeStruct(x_real.shape, jnp.complex64)
    )(x_real, x_imag, y_real, y_imag)

array1 = jnp.array([2+3j, 1-1j, 5+1j])
array2 = jnp.array([1+2j, 1-2j, 2+1j])

cdot_vectors(array1.real, array1.imag, array2.real, array2.imag)

One line in the error Traceback stuck to me:

NotImplementedError: cannot cast Value(%33 = "arith.addf"(%31, %32) <{fastmath = #arith.fastmath<none>}> : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>) to tensor<3xcomplex<f32>>

Any ideas what this happening? Are there any constraints on complex numbers or is it possible to have a kernel with 2 outputs?

[sharadmv]

I don't think this path is well trodden. My suggestion (for now) is to use separate inputs/outputs for complex values.

[bsaoptima]

So I managed to make the code work by having my output array / Refbeing 2D where I have a row for real values and another one for imaginary. Then I combine both outside of the kernel. I don't know if I would have gained some speedup by having that step in the kernel.

Something that I am curious about, can you have 2 outputs in one kernel?

[sharadmv]

Yes you can, out_shape can be a list or tuple.

[bsaoptima]

No what I meant was to have something like output_ref_1 and output_ref_2. Apologies for the lack of clarity.

[sharadmv]

Yes, so if you say out_shape=(jax.ShapeDtypeStruct(...), jax.ShapeDtypeStruct(...)), your kernel will be passed in two output references:

def kernel(..., out_ref1, out_ref2):
  ...

[bsaoptima]

Wow okay that is nice, thank you!