RUN pip install jax[cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION]==$JAX_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
/tmp/clean-layer.sh
{{ else }}
RUN pip install jax[cpu]==$JAX_VERSION && \
/tmp/clean-layer.sh
{{ end }}
# Install mxnet
{{ if eq .Accelerator "gpu" }}
RUN pip install mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \
/tmp/clean-layer.sh
{{ else }}
RUN pip install mxnet && \
/tmp/clean-layer.sh
{{ end}}
# Install GPU specific packages
{{ if eq .Accelerator "gpu" }}
# Install GPU-only packages
RUN pip install pycuda && \
pip install pynvrtc && \
# b/190622765 latest version is causing issue. nnabla fixed it in https://github.com/sony/nnabla/issues/892, waiting for new release before we can remove this pin.
# b/183041606#comment5: the Kaggle data proxy doesn't support these APIs. If the library is missing, it falls back to using a regular BigQuery query to fetch data.
# Stop Matplotlib printing junk to the console on first load
sed -i "s/^.*Matplotlib is building the font cache using fc-list.*$/# Warning removed by Kaggle/g" /opt/conda/lib/python3.7/site-packages/matplotlib/font_manager.py && \
# Make matplotlib output in Jupyter notebooks display correctly