4D conv_general_dilated gives XlaRuntimeError: INTERNAL: LHLO opcode convolution is not supported

[jajfi]

I'd like to use 4D convolution. This minimal example

import jax
key1, key2 = jax.random.split(jax.random.PRNGKey(0))
x = jax.random.normal(key1,(8, 11, 11, 11, 11, 1))
kernel = jax.random.normal(key2,(11, 11, 11, 11, 1, 3))
dn = jax.lax.conv_dimension_numbers(x.shape,kernel.shape,(("NXYZWC", "XYZWIO", "NXYZWC")))
output = jax.lax.conv_general_dilated(x,kernel,(1, 1, 1, 1),"SAME",dimension_numbers=dn)

gives the following error:

%convolution.3 = f32[8,11,11,11,11,3]{5,4,3,2,1,0} convolution(f32[8,11,11,11,11,1]{5,4,3,2,1,0} %Arg_0.1, f32[11,11,11,11,1,3]{5,4,3,2,1,0} %Arg_1.2), window={size=11x11x11x11 pad=5_5x5_5x5_5x5_5}, dim_labels=b0123f_0123io->b0123f, metadata={op_name="jit(conv_general_dilated)/jit(main)/conv_general_dilated[window_strides=(1, 1, 1, 1) padding=((5, 5), (5, 5), (5, 5), (5, 5)) lhs_dilation=(1, 1, 1, 1) rhs_dilation=(1, 1, 1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 5, 1, 2, 3, 4), rhs_spec=(5, 4, 0, 1, 2, 3), out_spec=(0, 5, 1, 2, 3, 4)) feature_group_count=1 batch_group_count=1 lhs_shape=(8, 11, 11, 11, 11, 1) rhs_shape=(11, 11, 11, 11, 1, 3) precision=None preferred_element_type=None]" source_file="/home/jajfi/Documents/Programming/jax/4dconv.py" source_line=6}Traceback (most recent call last):
  File "/home/jajfi/Documents/Programming/jax/4dconv.py", line 6, in <module>
    output = jax.lax.conv_general_dilated(x,kernel,(1, 1, 1, 1),"SAME",dimension_numbers=dn)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/lax/convolution.py", line 156, in conv_general_dilated
    return conv_general_dilated_p.bind(
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/core.py", line 329, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/core.py", line 332, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/core.py", line 712, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 118, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/util.py", line 254, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/util.py", line 247, in cached
    return f(*args, **kwargs)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 202, in xla_primitive_callable
    compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 359, in _xla_callable_uncached
    return computation.compile(_allow_propagation_to_outputs=True).unsafe_call
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py", line 3202, in compile
    executable = self._compile_unloaded(
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py", line 3170, in _compile_unloaded
    return UnloadedMeshExecutable.from_hlo(
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py", line 3439, in from_hlo
    xla_executable = dispatch.compile_or_get_cached(
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 1079, in compile_or_get_cached
    return backend_compile(backend, serialized_computation, compile_options,
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/jajfi/miniconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 1014, in backend_compile
    return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: LHLO opcode convolution is not supported.

If I change the example to 3D:

import jax
key1, key2 = jax.random.split(jax.random.PRNGKey(0))
x = jax.random.normal(key1,(8, 11, 11, 11, 1))
kernel = jax.random.normal(key2,(11, 11, 11, 1, 3))
dn = jax.lax.conv_dimension_numbers(x.shape,kernel.shape,(("NXYZC", "XYZIO", "NXYZC")))
output = jax.lax.conv_general_dilated(x,kernel,(1, 1, 1),"SAME",dimension_numbers=dn)
print(output.shape)

then I get the correct output.

I can reproduce this both on a fresh conda env with conda install jax cuda-nvcc -c conda-forge -c nvidia:

# Name                    Version                   Build  Channel
cuda-nvcc                 12.0.76                       0    nvidia
jax                       0.4.1              pyhd8ed1ab_0    conda-forge
jaxlib                    0.4.1           cuda112py310hf0bc174_201    conda-forge

and in a Docker container with this Dockerfile:

FROM nvcr.io/nvidia/tensorflow:22.12-tf2-py3
RUN pip install --upgrade pip
RUN pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Am I doing something wrong or this is a bug?