FROM wtongze/rocm-base:5.7.3
ARG TENSORFLOW_BRANCH="r2.11-rocm-enhanced"
ARG NUMPY_VERSION="1.22.4"
ARG DEBIAN_FRONTEND="noninteractive"

RUN sudo apt-get update && \
sudo apt-get install -y rocm-libs && \
sudo rm -rf /var/lib/apt/lists/*

WORKDIR /home/rocm

# Add GPU target
RUN bash -c 'echo -e "gfx803\ngfx900" | sudo tee /opt/rocm/bin/target.lst'

# Install bazelisk
RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.21.0/bazelisk-linux-amd64 && sudo chmod 755 /usr/local/bin/bazel

RUN git clone --depth 1 https://github.com/ROCm/tensorflow-upstream.git -b ${TENSORFLOW_BRANCH} tensorflow
WORKDIR /home/rocm/tensorflow

# > https://github.com/ROCm/ROCm/issues/1880
ENV PATH=$PATH:/home/rocm/.local/bin

RUN sudo apt-get update && \
sudo apt-get install -y python3-dev && \
sudo rm -rf /var/lib/apt/lists/*

RUN pip install -U --user pip numpy==${NUMPY_VERSION} wheel packaging requests opt_einsum patchelf
RUN pip install -U --user keras_preprocessing --no-deps

RUN sudo ln -s /usr/bin/python3 /usr/bin/python
RUN ./build_rocm_python3 /opt/rocm
RUN cp -a /tmp/tensorflow_pkg /home/rocm

CMD [ "/bin/bash", "-l" ]