x == y == z

[cool-RR]

How does one run something like x == y == z on JAX arrays? I want to get an array of the same shape that has True only when items from all three arrays are equal. This fails because x == y just evaluates to booleans which are then compared to z.

[jakevdp]

You can do this:

(x == y) & (y == z)

Note that the parentheses are required here due to operator precedence rules.

This fails because x == y just evaluates to booleans which are then compared to z.

Side note here, but it fails for a different reason: Python treats operator chaining like x == y == z as equivalent to (x == y) and (y == z), and the and keyword will try to cast the things on either side to bool, and bool(x == y) fails.

[cool-RR]

I see. That's a little sad but I can live with it.

[cool-RR]

I implemented this now:

def are_equal(arrays: Iterable[jax.Array]) -> jax.Array:
    processed_arrays = jnp.broadcast_arrays(*arrays)
    boolean_arrays = [processed_array == processed_arrays[0]
                      for processed_array in processed_arrays[1:]]
    return jnp.all(jnp.stack(boolean_arrays), axis=0)

Probably could be made more efficient, but helps me not repeat y in the expression above.

[jakevdp]

Seems like a lot of code to write in order to avoid typing y twice, but other than that it looks fine.