jax

c061fd7b - improve error for unimplemented parallel primitive