How would I get the cartesian product of two arrays in Jax?