jnp.vectorize static argument

[chrisflesher]

Hello,

I like the syntactic sugar of using jnp.vectorize but sometimes have situations where I can't use it. See the example below:

@functools.partial(jnp.vectorize, signature=('(a,b),(a,b)->(c,d)'), excluded=(2,))
def _cross_correlate(source_image: jax.Array,
                     target_image: jax.Array,
                     ifft: typing.Callable = lambda x: jnp.real(jnp.fft.ifft2(x)),
                     ) -> jax.Array:
    fft = jnp.fft.fft2
    return ifft(fft(source_image) * fft(target_image[::-1, ::-1])

But when I run this it says:

ValueError: excluded=(2,) is invalid for 2 argument(s)

Is there a way to set a default value for ifft without a lot of extra boilerplate?

[jakevdp]

No, unfortunately jnp.vectorize only supports positional arguments for exclusion. Until jnp.vectorize supports excluding elements by keyword, the best option will probably be to write this as a wrapper function over positional args; something like this:

def _cross_correlate(source_image: jax.Array,
                     target_image: jax.Array,
                     ifft: typing.Callable = lambda x: jnp.real(jnp.fft.ifft2(x)),
                     ) -> jax.Array:
  return _cross_correlate_impl(source_image, target_image, iff)

@functools.partial(jnp.vectorize, signature=('(a,b),(a,b)->(c,d)'), excluded=(2,))
def _cross_correlate_impl(source_image: jax.Array, target_image: jax.Array, ifft: typing.Callable) -> jax.Array:
    fft = jnp.fft.fft2
    return ifft(fft(source_image) * fft(target_image[::-1, ::-1])

[jakevdp]

I opened #19066 to track adding this feature. Once that is complete, I believe passing excluded=(2, "ifft") to your original function would be sufficient to get the behavior you want.

[chrisflesher]

Cool, thanks!

[jakevdp]

#19072 should address this.