Issue 1:
The following operations would run easily in both Tensorflow and PyTorch (GPU), but would cause OOM in Jax. When it runs on PyTorch, GPU memory usage is only ~2GB.

Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

import jax.numpy as jnp
import jax.random as jrandom
A = jrandom.normal(jrandom.PRNGKey(0), [300000, 50])
print(A.shape)
jnp.linalg.svd(A, full_matrices=False)

If applicable, include full error messages/tracebacks.

---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
/tmp/ipykernel_57631/2574565684.py in <module>
3 A = jrandom.normal(jrandom.PRNGKey(0), [300000, 50])
4 print(A.shape)
----> 5 jnp.linalg.svd(A, full_matrices=False)

[... skipping hidden 6 frame]

/opt/conda/lib/python3.7/site-packages/jax/_src/dispatch.py in _execute_compiled(name, compiled, output_buffer_counts, result_handlers, kept_var_idx, *args)
442 input_bufs = util.flatten(
443 device_put(x, device) for i, x in enumerate(args) if i in kept_var_idx)
--> 444 out_bufs = compiled.execute(input_bufs)
445 check_special(name, out_bufs)
446 if output_buffer_counts is None:

RuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 360336162820 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
parameter allocation: 57.22MiB
constant allocation: 0B
maybe_live_out allocation: 57.23MiB
preallocated temp allocation: 335.59GiB
preallocated temp fragmentation: 112B (0.00%)
total allocation: 335.70GiB
total fragmentation: 10.0KiB (0.00%)
Peak buffers:
Buffer 1:
Size: 335.28GiB
XLA Label: custom-call
Shape: f32[300000,300000]
==========================

Buffer 2:
Size: 263.36MiB
XLA Label: custom-call
Shape: f32[69038144]
==========================

Buffer 3:
Size: 57.22MiB
XLA Label: custom-call
Shape: f32[300000,50]
==========================

Buffer 4:
Size: 57.22MiB
Entry Parameter Subshape: f32[300000,50]
==========================

Buffer 5:
Size: 57.22MiB
Operator: op_name="jit(svd)/jit(main)/svd[full_matrices=False compute_uv=True]" source_file="/tmp/ipykernel_57631/2574565684.py" source_line=5
XLA Label: fusion
Shape: f32[300000,50]
==========================

Buffer 6:
Size: 9.8KiB
XLA Label: custom-call
Shape: f32[50,50]
==========================

Buffer 7:
Size: 9.8KiB
Operator: op_name="jit(svd)/jit(main)/svd[full_matrices=False compute_uv=True]" source_file="/tmp/ipykernel_57631/2574565684.py" source_line=5
XLA Label: fusion
Shape: f32[50,50]
==========================

Buffer 8:
Size: 200B
Operator: op_name="jit(svd)/jit(main)/svd[full_matrices=False compute_uv=True]" source_file="/tmp/ipykernel_57631/2574565684.py" source_line=5
XLA Label: fusion
Shape: f32[50]
==========================

Buffer 9:
Size: 48B
XLA Label: custom-call
Shape: (f32[300000,50], f32[50], f32[300000,300000], f32[50,50], s32[], /*index=5*/f32[69038144])
==========================

Buffer 10:
Size: 24B
XLA Label: tuple
Shape: (f32[300000,50], f32[50], f32[50,50])
==========================

Buffer 11:
Size: 4B
XLA Label: custom-call
Shape: s32[]
==========================

Issue 2:
There is at least one other report of a build failure: #11597.

In conda-forge, our jaxlib builds completely failed for different reasons after trying to update to jaxlib v0.3.15. They fail differently on cuda and non-cuda builds. For cuda, which is the more pertinent error, it looks like:

INFO: Found 1 target...
[0 / 4] [Prepa] Creating source manifest for //build:build_wheel
[27 / 143] Compiling src/google/protobuf/stubs/structurally_valid.cc; 0s local ... (2 actions, 1 running)
[38 / 143] Compiling src/google/protobuf/generated_message_util.cc; 1s local ... (2 actions running)
ERROR: /home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/share/bazel/8372a0709531b2dba2d9a9ce09c7583d/external/com_google_protobuf/BUILD:111:11: Linking external/com_google_protobuf/libprotobuf_lite.lo failed: (Exit 1): ar failed: error executing command
(cd /home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/share/bazel/8372a0709531b2dba2d9a9ce09c7583d/execroot/__main__ && \
exec env - \
CUDA_TOOLKIT_PATH=/usr/local/cuda \
CUDNN_INSTALL_PATH=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac \
GCC_HOST_COMPILER_PATH=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin/x86_64-conda-linux-gnu-gcc \
GCC_HOST_COMPILER_PREFIX=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin \
PATH=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin:/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac/bin:/opt/conda/condabin:/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env:/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin:/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac:/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac/bin:/opt/conda/bin:/opt/conda/condabin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/home/conda/bin:/usr/local/cuda/bin \
PWD=/proc/self/cwd \
TF_CUDA_COMPUTE_CAPABILITIES=sm_35,sm_50,sm_60,sm_62,sm_70,sm_72,sm_75,sm_80,sm_86,compute_86 \
TF_CUDA_PATHS=/usr/local/cuda,/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac \
TF_CUDA_VERSION=11.2 \
TF_CUDNN_VERSION=8 \
/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin/ar @bazel-out/k8-opt/bin/external/com_google_protobuf/libprotobuf_lite.lo-2.params)
# Configuration: 740083763ed3206026a09abc2faa9ac5a84e6f9dfcfd22c27729d3542ad1c4a3
# Execution platform: @local_execution_config_platform//:platform
src/main/tools/process-wrapper-legacy.cc:80: "execvp(/home/conda/feedstock_root/build_artifacts/jaxlib_1658602566191/_build_env/bin/ar, ...)": No such file or directory
Target //build:build_wheel failed to build
INFO: Elapsed time: 56.913s, Critical Path: 3.40s
INFO: 47 processes: 18 internal, 29 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully

On non-cuda, we get this error, which I think is more likely something on our end, but I cannot figure it out yet...

[6,508 / 6,659] Compiling llvm/lib/Target/X86/X86ISelLowering.cpp; 35s local ... (2 actions running)
Target //build:build_wheel up-to-date:
bazel-bin/build/build_wheel
INFO: Elapsed time: 13046.285s, Critical Path: 205.49s
INFO: 7181 processes: 616 internal, 6565 local.
INFO: Build completed successfully, 7181 total actions
INFO: Running command line: bazel-bin/build/build_wheel '--output_path=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/work/dist' '--cpu=x86_64'
INFO: Build completed successfully, 7181 total actions
/home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac/lib/python3.8/site-packages/setuptools/command/install.py:34: SetuptoolsDeprecationWarning: setup.py install is deprecated. Use build and pip and other standards-based tools.
warnings.warn(
Output wheel: /home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/work/dist/jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl

To install the newly-built jaxlib wheel, run:
pip install /home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/work/dist/jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl

INFO: Reading rc options for 'shutdown' from /home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/work/.jax_configure.bazelrc:
Inherited 'common' options: --crosstool_top=//bazel_toolchain:toolchain --logging=6 --verbose_failures --toolchain_resolution_debug --define=PREFIX=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac --define=PROTOBUF_INCLUDE_PATH=/home/conda/feedstock_root/build_artifacts/jaxlib_1658602489386/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac/include --local_cpu_resources=2 --cpu k8
ERROR: --crosstool_top=//bazel_toolchain:toolchain :: Unrecognized option: --crosstool_top=//bazel_toolchain:toolchain
b''

Could the changelog be updated at some point? Note that at least one more PR seems to be destined for v0.3.15 has been added, even though v0.3.15 has been released already on pypi... #11602

Could someone have a look at the above errors and see if something obvious we should be aware of in packaging at conda-forge? As a reminder we have our own custom bazel toolchain, so most likely at least some of the errors are related to that. I am still working on determining what exactly. I only recently got involved in packaging using bazel on conda-forge...

Current live PR: conda-forge/jaxlib-feedstock#118

Previous PR with failures: conda-forge/jaxlib-feedstock#116

Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

import jax.numpy as jnp
print(jnp.arange(10))
# [0 1 2 3 4 5 6 7 8 9]

If applicable, include full error messages/tracebacks.