How to use jax.eval_shape with lax.conv_general_dilated?

[adonath]

I have tried to evaluate the shape of lax.conv_general_dilated operation using jax.eval_shape and an abstract ShapeDtypeStruct. However both of the naive approaches fail:

import jax
from jax import lax
from functools import partial

x = jax.ShapeDtypeStruct((1, 1, 28, 28), jnp.float32)
kernel = jax.ShapeDtypeStruct((32, 1, 3, 3), jnp.float32)

# does not work
jax.eval_shape(lax.conv_general_dilated, x, kernel, (1, 1), "SAME")

# use a custom jvp to declare which args are non-diff
@partial(jax.custom_jvp, nondiff_argnums=(2, 3))
def convolve_wrap(x, kernel, stride, padding):
    """Convolve 2D."""
    return lax.conv_general_dilated(
        x,
        kernel,
        window_strides=(stride, stride),
        padding=padding,
    )

# does not work either
jax.eval_shape(convolve_wrap, x, kernel, (1, 1), "SAME")

Both fail with:

TypeError: Argument 'SAME' of type <class 'str'> is not a valid JAX type

I'd be happy for any hint on how to use it correctly. Thanks!

[jakevdp]

All arguments of jax.eval_shape will be treated as dynamic; any static arguments should be bound to the function using a closure or a partial. For example:

jax.eval_shape(
  partial(lax.conv_general_dilated, window_strides=(1, 1), padding="SAME"),
  x, kernel)

[adonath]

Perfect, thanks for the explanation!