<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">I use singularity to build an image and follow the README.md using the command like below.</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="./build_rocm.sh"><pre class="notranslate">./build_rocm.sh</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.<br>
these are all the outputs.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Singularity&gt; ./build_rocm.sh
+ ROCM_TF_FORK_REPO=https://github.com/ROCmSoftwarePlatform/tensorflow-upstream
+ ROCM_TF_FORK_BRANCH=develop-upstream
+ ROCM_PATH=/opt/rocm/
+ python3 ../build.py --enable_rocm --rocm_path=/opt/rocm/ --bazel_options=--override_repository=org_tensorflow=/tmp/tensorflow-upstream

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: ./bazel-5.1.0-linux-x86_64
Bazel version: 5.1.0
Python binary path: /opt/conda/bin/python3
Python version: 3.9
NumPy version: 1.20.3
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: yes
ROCm toolkit path: /opt/rocm/
ROCm amdgpu targets: gfx900,gfx906,gfx908,gfx90a,gfx1030

Building XLA and installing it in the jaxlib source tree...
./bazel-5.1.0-linux-x86_64 run --verbose_failures=true --override_repository=org_tensorflow=/tmp/tensorflow-upstream --config=avx_posix --config=mkl_open_source_only --config=rocm :build_wheel -- --output_path=/home/code/jax/build/rocm/dist --cpu=x86_64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /home/code/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /home/code/jax/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --experimental_cc_shared_library --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /home/code/jax/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/opt/conda/bin/python3 --action_env=PYENV_ROOT --python_path=/opt/conda/bin/python3 --action_env ROCM_PATH=/opt/rocm/ --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /home/code/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /home/code/jax/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /home/code/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:rocm in file /home/code/jax/.bazelrc: --crosstool_top=@local_config_rocm//crosstool:toolchain --define=using_rocm=true --define=using_rocm_hipcc=true --define=xla_python_enable_gpu=true --repo_env TF_NEED_ROCM=1 --action_env TF_ROCM_AMDGPU_TARGETS=gfx900,gfx906,gfx908
INFO: Found applicable config definition build:rocm in file /home/code/jax/.jax_configure.bazelrc: --action_env TF_ROCM_AMDGPU_TARGETS=gfx900,gfx906,gfx908,gfx90a,gfx1030
INFO: Found applicable config definition build:linux in file /home/code/jax/.bazelrc: --config=posix --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter
INFO: Found applicable config definition build:posix in file /home/code/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading:
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/tensorflow/runtime/archive/2123408fb43a5c4afdf87dafd67117d9c0ff70cd.tar.gz failed: class java.io.FileNotFoundException GET returned 404 Not Found
Loading: 0 packages loaded
Loading: 0 packages loaded
INFO: Repository local_config_rocm instantiated at:
  /home/code/jax/WORKSPACE:31:14: in &lt;toplevel&gt;
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace2.bzl:869:19: in workspace
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace2.bzl:101:19: in _tf_toolchains
Repository rule rocm_configure defined at:
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl:843:33: in &lt;toplevel&gt;
ERROR: An error occurred during the fetch of repository 'local_config_rocm':
   Traceback (most recent call last):
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 824, column 38, in _rocm_autoconf_impl
                _create_local_rocm_repository(repository_ctx)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 550, column 35, in _create_local_rocm_repository
                rocm_config = _get_rocm_config(repository_ctx, bash_bin, find_rocm_config_script)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 398, column 30, in _get_rocm_config
                config = find_rocm_config(repository_ctx, find_rocm_config_script)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 376, column 41, in find_rocm_config
                exec_result = _exec_find_rocm_config(repository_ctx, script_path)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 372, column 19, in _exec_find_rocm_config
                return execute(repository_ctx, [python_bin, &quot;-c&quot;, decompress_and_execute_cmd])
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/remote_config/common.bzl&quot;, line 230, column 13, in execute
                fail(
Error in fail: Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
ERROR: /home/code/jax/WORKSPACE:31:14: fetching rocm_configure rule //external:local_config_rocm: Traceback (most recent call last):
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 824, column 38, in _rocm_autoconf_impl
                _create_local_rocm_repository(repository_ctx)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 550, column 35, in _create_local_rocm_repository
                rocm_config = _get_rocm_config(repository_ctx, bash_bin, find_rocm_config_script)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 398, column 30, in _get_rocm_config
                config = find_rocm_config(repository_ctx, find_rocm_config_script)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 376, column 41, in find_rocm_config
                exec_result = _exec_find_rocm_config(repository_ctx, script_path)
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl&quot;, line 372, column 19, in _exec_find_rocm_config
                return execute(repository_ctx, [python_bin, &quot;-c&quot;, decompress_and_execute_cmd])
        File &quot;/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/remote_config/common.bzl&quot;, line 230, column 13, in execute
                fail(
Error in fail: Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
INFO: Repository bazel_skylib instantiated at:
  /home/code/jax/WORKSPACE:28:14: in &lt;toplevel&gt;
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace3.bzl:21:17: in workspace
Repository rule http_archive defined at:
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/bazel_tools/tools/build_defs/repo/http.bzl:353:31: in &lt;toplevel&gt;
ERROR: Skipping ':build_wheel': no such package '@local_config_rocm//rocm': Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
WARNING: Target pattern parsing failed.
ERROR: no such package '@local_config_rocm//rocm': Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
INFO: Elapsed time: 77.870s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (0 packages loaded)
b''
Traceback (most recent call last):
  File &quot;/home/code/jax/build/rocm/../build.py&quot;, line 527, in &lt;module&gt;
    main()
  File &quot;/home/code/jax/build/rocm/../build.py&quot;, line 522, in main
    shell(command)
  File &quot;/home/code/jax/build/rocm/../build.py&quot;, line 53, in shell
    output = subprocess.check_output(cmd)
  File &quot;/opt/conda/lib/python3.9/subprocess.py&quot;, line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File &quot;/opt/conda/lib/python3.9/subprocess.py&quot;, line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.0-linux-x86_64', 'run', '--verbose_failures=true', '--override_repository=org_tensorflow=/tmp/tensorflow-upstream', '--config=avx_posix', '--config=mkl_open_source_only', '--config=rocm', ':build_wheel', '--', '--output_path=/home/code/jax/build/rocm/dist', '--cpu=x86_64']' returned non-zero exit status 1."><pre class="notranslate"><code class="notranslate">Singularity&gt; ./build_rocm.sh
+ ROCM_TF_FORK_REPO=https://github.com/ROCmSoftwarePlatform/tensorflow-upstream
+ ROCM_TF_FORK_BRANCH=develop-upstream
+ ROCM_PATH=/opt/rocm/
+ python3 ../build.py --enable_rocm --rocm_path=/opt/rocm/ --bazel_options=--override_repository=org_tensorflow=/tmp/tensorflow-upstream

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: ./bazel-5.1.0-linux-x86_64
Bazel version: 5.1.0
Python binary path: /opt/conda/bin/python3
Python version: 3.9
NumPy version: 1.20.3
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: yes
ROCm toolkit path: /opt/rocm/
ROCm amdgpu targets: gfx900,gfx906,gfx908,gfx90a,gfx1030

Building XLA and installing it in the jaxlib source tree...
./bazel-5.1.0-linux-x86_64 run --verbose_failures=true --override_repository=org_tensorflow=/tmp/tensorflow-upstream --config=avx_posix --config=mkl_open_source_only --config=rocm :build_wheel -- --output_path=/home/code/jax/build/rocm/dist --cpu=x86_64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /home/code/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /home/code/jax/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --experimental_cc_shared_library --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /home/code/jax/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/opt/conda/bin/python3 --action_env=PYENV_ROOT --python_path=/opt/conda/bin/python3 --action_env ROCM_PATH=/opt/rocm/ --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /home/code/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /home/code/jax/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /home/code/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:rocm in file /home/code/jax/.bazelrc: --crosstool_top=@local_config_rocm//crosstool:toolchain --define=using_rocm=true --define=using_rocm_hipcc=true --define=xla_python_enable_gpu=true --repo_env TF_NEED_ROCM=1 --action_env TF_ROCM_AMDGPU_TARGETS=gfx900,gfx906,gfx908
INFO: Found applicable config definition build:rocm in file /home/code/jax/.jax_configure.bazelrc: --action_env TF_ROCM_AMDGPU_TARGETS=gfx900,gfx906,gfx908,gfx90a,gfx1030
INFO: Found applicable config definition build:linux in file /home/code/jax/.bazelrc: --config=posix --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter
INFO: Found applicable config definition build:posix in file /home/code/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading:
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/tensorflow/runtime/archive/2123408fb43a5c4afdf87dafd67117d9c0ff70cd.tar.gz failed: class java.io.FileNotFoundException GET returned 404 Not Found
Loading: 0 packages loaded
Loading: 0 packages loaded
INFO: Repository local_config_rocm instantiated at:
  /home/code/jax/WORKSPACE:31:14: in &lt;toplevel&gt;
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace2.bzl:869:19: in workspace
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace2.bzl:101:19: in _tf_toolchains
Repository rule rocm_configure defined at:
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl:843:33: in &lt;toplevel&gt;
ERROR: An error occurred during the fetch of repository 'local_config_rocm':
   Traceback (most recent call last):
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 824, column 38, in _rocm_autoconf_impl
                _create_local_rocm_repository(repository_ctx)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 550, column 35, in _create_local_rocm_repository
                rocm_config = _get_rocm_config(repository_ctx, bash_bin, find_rocm_config_script)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 398, column 30, in _get_rocm_config
                config = find_rocm_config(repository_ctx, find_rocm_config_script)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 376, column 41, in find_rocm_config
                exec_result = _exec_find_rocm_config(repository_ctx, script_path)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 372, column 19, in _exec_find_rocm_config
                return execute(repository_ctx, [python_bin, "-c", decompress_and_execute_cmd])
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/remote_config/common.bzl", line 230, column 13, in execute
                fail(
Error in fail: Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
ERROR: /home/code/jax/WORKSPACE:31:14: fetching rocm_configure rule //external:local_config_rocm: Traceback (most recent call last):
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 824, column 38, in _rocm_autoconf_impl
                _create_local_rocm_repository(repository_ctx)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 550, column 35, in _create_local_rocm_repository
                rocm_config = _get_rocm_config(repository_ctx, bash_bin, find_rocm_config_script)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 398, column 30, in _get_rocm_config
                config = find_rocm_config(repository_ctx, find_rocm_config_script)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 376, column 41, in find_rocm_config
                exec_result = _exec_find_rocm_config(repository_ctx, script_path)
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/gpus/rocm_configure.bzl", line 372, column 19, in _exec_find_rocm_config
                return execute(repository_ctx, [python_bin, "-c", decompress_and_execute_cmd])
        File "/root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/third_party/remote_config/common.bzl", line 230, column 13, in execute
                fail(
Error in fail: Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
INFO: Repository bazel_skylib instantiated at:
  /home/code/jax/WORKSPACE:28:14: in &lt;toplevel&gt;
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/org_tensorflow/tensorflow/workspace3.bzl:21:17: in workspace
Repository rule http_archive defined at:
  /root/.cache/bazel/_bazel_root/7777a22c05c38dcb5674638712ab6fae/external/bazel_tools/tools/build_defs/repo/http.bzl:353:31: in &lt;toplevel&gt;
ERROR: Skipping ':build_wheel': no such package '@local_config_rocm//rocm': Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
WARNING: Target pattern parsing failed.
ERROR: no such package '@local_config_rocm//rocm': Repository command failed
ERROR: ROCm version file not found in ['include/rocm-core/rocm_version.h', 'include/rocm_version.h']
INFO: Elapsed time: 77.870s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (0 packages loaded)
b''
Traceback (most recent call last):
  File "/home/code/jax/build/rocm/../build.py", line 527, in &lt;module&gt;
    main()
  File "/home/code/jax/build/rocm/../build.py", line 522, in main
    shell(command)
  File "/home/code/jax/build/rocm/../build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/opt/conda/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/opt/conda/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.0-linux-x86_64', 'run', '--verbose_failures=true', '--override_repository=org_tensorflow=/tmp/tensorflow-upstream', '--config=avx_posix', '--config=mkl_open_source_only', '--config=rocm', ':build_wheel', '--', '--output_path=/home/code/jax/build/rocm/dist', '--cpu=x86_64']' returned non-zero exit status 1.
</code></pre></div>
<p dir="auto">Hello,<br>
After discussion  <a href="https://github.com/google/jax/discussions/15204" data-hovercard-type="discussion" data-hovercard-url="/google/jax/discussions/15204/hovercard">#15204</a></p>
<p dir="auto">Here is a code to deal with the convolution using FFT of two 2D real arrays. It is mostly inspired from the <code class="notranslate">scipy.fftconvolve</code>function.  One use-case &amp; test of this function is available on Colab  <a href="https://colab.research.google.com/drive/1Ze7nxlQQc05FT_bt1YuxX7PjjYBmBe17?usp=sharing" rel="nofollow">HERE</a> with the JAX 0.4.6 version (this is the current one at least today). For sure my code is not the best one. Feel free to improve it.</p>
<ul dir="auto">
<li>I do nit know if the second array in2 can be smaller in size than the first one</li>
<li>for sure the complex-complex case would need some adaptation</li>
<li>finally the <code class="notranslate">next_fast_len</code> available in scipy  would improve the computation of the "good" FFT size.</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@jax.jit
def convolve2r(in1,in2):

  # suppose
  # in1 = jnp.asarray(x1)
  # in2 = jnp.asarray(x2)
  #  &quot;assert&quot; in1.ndim == in2.ndim == 2
  #  &quot;assert&quot; in1.size == in2.size
  #  &quot;assert&quot; in1.dtype.kind == 'f' and in2.dtype.kind == &quot;f&quot;

  axes = range(in1.ndim)

  s1 = in1.shape
  s2 = in2.shape
  shape = [s1[i] + s2[i] -1 for i in range(in1.ndim)]
  
  # crude next_fast_len equivalent with just power of 2
  fshape = 2**np.ceil(np.log2(np.array(shape))).astype('uint16')

  fslice = tuple([slice(sz) for sz in shape])

  sp1 = jnp.fft.fftn(in1, fshape, axes=axes)
  sp2 = jnp.fft.fftn(in2, fshape, axes=axes)
  conv = jnp.fft.irfftn(sp1 * sp2, fshape, axes=axes)
  conv = conv[tuple(map(slice, shape))]   #conv = jax.lax.slice(conv,(0,0),(shape))

  def _centered(arr, newshape):
    # Return the center newshape portion of the array.
    newshape = np.asarray(newshape)
    currshape = np.array(arr.shape)
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    return arr[tuple(myslice)]

  return _centered(conv, s1)"><pre class="notranslate"><span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">convolve2r</span>(<span class="pl-s1">in1</span>,<span class="pl-s1">in2</span>):

  <span class="pl-c"># suppose</span>
  <span class="pl-c"># in1 = jnp.asarray(x1)</span>
  <span class="pl-c"># in2 = jnp.asarray(x2)</span>
  <span class="pl-c">#  "assert" in1.ndim == in2.ndim == 2</span>
  <span class="pl-c">#  "assert" in1.size == in2.size</span>
  <span class="pl-c">#  "assert" in1.dtype.kind == 'f' and in2.dtype.kind == "f"</span>

  <span class="pl-s1">axes</span> <span class="pl-c1">=</span> <span class="pl-en">range</span>(<span class="pl-s1">in1</span>.<span class="pl-s1">ndim</span>)

  <span class="pl-s1">s1</span> <span class="pl-c1">=</span> <span class="pl-s1">in1</span>.<span class="pl-s1">shape</span>
  <span class="pl-s1">s2</span> <span class="pl-c1">=</span> <span class="pl-s1">in2</span>.<span class="pl-s1">shape</span>
  <span class="pl-s1">shape</span> <span class="pl-c1">=</span> [<span class="pl-s1">s1</span>[<span class="pl-s1">i</span>] <span class="pl-c1">+</span> <span class="pl-s1">s2</span>[<span class="pl-s1">i</span>] <span class="pl-c1">-</span><span class="pl-c1">1</span> <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">in1</span>.<span class="pl-s1">ndim</span>)]
  
  <span class="pl-c"># crude next_fast_len equivalent with just power of 2</span>
  <span class="pl-s1">fshape</span> <span class="pl-c1">=</span> <span class="pl-c1">2</span><span class="pl-c1">**</span><span class="pl-s1">np</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">np</span>.<span class="pl-en">log2</span>(<span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">shape</span>))).<span class="pl-en">astype</span>(<span class="pl-s">'uint16'</span>)

  <span class="pl-s1">fslice</span> <span class="pl-c1">=</span> <span class="pl-en">tuple</span>([<span class="pl-en">slice</span>(<span class="pl-s1">sz</span>) <span class="pl-k">for</span> <span class="pl-s1">sz</span> <span class="pl-c1">in</span> <span class="pl-s1">shape</span>])

  <span class="pl-s1">sp1</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">fft</span>.<span class="pl-en">fftn</span>(<span class="pl-s1">in1</span>, <span class="pl-s1">fshape</span>, <span class="pl-s1">axes</span><span class="pl-c1">=</span><span class="pl-s1">axes</span>)
  <span class="pl-s1">sp2</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">fft</span>.<span class="pl-en">fftn</span>(<span class="pl-s1">in2</span>, <span class="pl-s1">fshape</span>, <span class="pl-s1">axes</span><span class="pl-c1">=</span><span class="pl-s1">axes</span>)
  <span class="pl-s1">conv</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">fft</span>.<span class="pl-en">irfftn</span>(<span class="pl-s1">sp1</span> <span class="pl-c1">*</span> <span class="pl-s1">sp2</span>, <span class="pl-s1">fshape</span>, <span class="pl-s1">axes</span><span class="pl-c1">=</span><span class="pl-s1">axes</span>)
  <span class="pl-s1">conv</span> <span class="pl-c1">=</span> <span class="pl-s1">conv</span>[<span class="pl-en">tuple</span>(<span class="pl-en">map</span>(<span class="pl-s1">slice</span>, <span class="pl-s1">shape</span>))]   <span class="pl-c">#conv = jax.lax.slice(conv,(0,0),(shape))</span>

  <span class="pl-k">def</span> <span class="pl-en">_centered</span>(<span class="pl-s1">arr</span>, <span class="pl-s1">newshape</span>):
    <span class="pl-c"># Return the center newshape portion of the array.</span>
    <span class="pl-s1">newshape</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">asarray</span>(<span class="pl-s1">newshape</span>)
    <span class="pl-s1">currshape</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">arr</span>.<span class="pl-s1">shape</span>)
    <span class="pl-s1">startind</span> <span class="pl-c1">=</span> (<span class="pl-s1">currshape</span> <span class="pl-c1">-</span> <span class="pl-s1">newshape</span>) <span class="pl-c1">//</span> <span class="pl-c1">2</span>
    <span class="pl-s1">endind</span> <span class="pl-c1">=</span> <span class="pl-s1">startind</span> <span class="pl-c1">+</span> <span class="pl-s1">newshape</span>
    <span class="pl-s1">myslice</span> <span class="pl-c1">=</span> [<span class="pl-en">slice</span>(<span class="pl-s1">startind</span>[<span class="pl-s1">k</span>], <span class="pl-s1">endind</span>[<span class="pl-s1">k</span>]) <span class="pl-k">for</span> <span class="pl-s1">k</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-en">len</span>(<span class="pl-s1">endind</span>))]
    <span class="pl-k">return</span> <span class="pl-s1">arr</span>[<span class="pl-en">tuple</span>(<span class="pl-s1">myslice</span>)]

  <span class="pl-k">return</span> <span class="pl-en">_centered</span>(<span class="pl-s1">conv</span>, <span class="pl-s1">s1</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
</ul>