Is there a jax equivalent to numpy.mod.outer?

[romilly]

I'm looking for a jax expression that's the equivalent to 0 == numpy.mod.outer(a, b) where a, b are vectors.

[jakevdp]

Yes, for example you can do this via broadcasting within a standard mod operation:

import numpy as np
import jax.numpy as jnp

a = jnp.array([1, 2, 3, 4])
b = jnp.array([1, 2, 3, 4, 5])

print(a[:, None] % b)
# [[0 1 1 1 1]
#  [0 0 2 2 2]
#  [0 1 0 3 3]
#  [0 0 1 0 4]]

[romilly]

Thanks! Problem solved, and something important learned.