Export stablehlo from jax function with kwargs

[sechkova]

Hey there! Following this guide: https://openxla.org/stablehlo/tutorials/jax-export I am trying to export stablehlo from a function that takes additional arguments.

Below is an example with convolution.

What should be the correct type of kwargs (params in my example)?
Given the errors I get (some of them added as comments) params should be hash-able but also having shape and type but not only symbolic specs, the real value is needed.

Is this the correct way of exporting this kind of functions?
I tried using wrapping conv in a lambda and putting params in the body. This approach works but produces different stablehlo output.

@functools.partial(
    jax.jit,
    static_argnames=("strides"),
)
def conv(x, w, strides):
    return lax.conv(
      lhs=x,
      rhs=w,
      window_strides = strides,
      padding = 'SAME'
  )

def gen_conv():

    x = jnp.array([[[[0., 1., 2., 3., 4.],  # (1, 1, 5, 5) input tensor
                    [5., 6., 7., 8., 9.],
                    [10., 11., 12., 13., 14.],
                    [15., 16., 17., 18., 19.],
                    [20., 21., 22., 23., 24.]]]]).astype(np.float32)
    W = jnp.array([[[[1., 1., 1.],  # (1, 1, 3, 3) tensor for convolution weights
                    [1., 1., 1.],
                    [1., 1., 1.]]]]).astype(np.float32)
    strides = jnp.array([1,1])

    inputs = [x, W]
    params = {'strides' : jnp.array((1,1))}

    try:
        stablehlo_str = export.export(conv)(*inputs, **params).mlir_module()
    except ValueError as e:
        print(e)
        # Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses.
        # Static argument (name strides) of type <class 'jaxlib.xla_extension.ArrayImpl'> for function conv is non-hashable.
        pass

    try:
        shaped_params = export.symbolic_args_specs(params, polymorphic_shapes=None)
        stablehlo_str = export.export(conv)(*inputs, **shaped_params).mlir_module()
    except TypeError as e:
        print(e)
        # unsupported operand type(s) for //: 'int' and 'ShapeDtypeStruct'
        pass

if __name__ == '__main__':
    gen_conv()

[jakevdp]

Array arguments cannot be static, so you cannot mark strides as a static argument and pass jnp.array([1, 1]) to it.

To fix this you can either (1) not mark strides as static, or (2) use a static, hashable value such as params = {'strides': (1, 1)}

[sechkova]

Thanks for the answers, it makes sense but in this particular case of export, if strides is a tuple, I get an error about missing attribute shape:

AttributeError: 'int' object has no attribute 'shape'

And if strides is not static and I use array, then this error occurs:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape int32[2].
The error occurred while tracing the function conv for jit. This concrete value was not available in Python because it depends on the value of the argument strides.