.at[] with complex numbers

import jax
import jax.numpy as jnp

@jax.jit
def f():
    return jnp.zeros((1,), dtype=jnp.complex64).at[0].set(1)

f()

DeviceArray([3.6602185e-37+0.j], dtype=complex64)

[mariogeiger]

You are always so efficient to fix issues ❤️

[hawkinsp]

The MHLO canonicalizer is mishandling simplifications of complex scatters. This will need a fix in upstream MLIR most likely.