# JAX requirements, passed into container by defining the ARG 
# REQUIREMENTS_FILE=jax.requirements.txt


setuptools
wheel
cloudpickle
colorama>=0.4.4
matplotlib
pillow>=9.1.0
rich
absl-py
portpicker
six
opt-einsum
auditwheel
typing_extensions
importlib_metadata>=4.6
numpy==1.26.0;python_version=="3.12"
numpy==1.23.4;python_version=="3.11"
numpy==1.22.4;python_version<"3.11"
scipy==1.11.2;python_version=="3.12"

  else
    TAG="pr-${KOKORO_GITHUB_PULL_REQUEST_NUMBER}"
  fi
fi

# Build for both JAX and TF usage.  We do these in one place because they share
# almost all of the same cache layers
export DOCKER_BUILDKIT=1
for target in jax tf; do
  IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
  docker pull "$IMAGE" || true

RUN echo $CACHEBUSTER

# Setup build and environment
COPY devel.usertools /usertools
COPY devel.bashrc /root/.bashrc
COPY ld.so.conf /dt10/etc/

# Setup JAX Python environment.
FROM devel as jax
RUN /setup.packages.sh /cuda.packages.txt
COPY jax.requirements.txt /devel.requirements.txt
COPY setup.python.sh /setup.python.sh
RUN /setup.python.sh python3.9 devel.requirements.txt

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

keras-nightly ~= 3.0.0.dev
tb-nightly ~= 2.17.0.a

# Test dependencies
grpcio >= 1.24.3, < 2.0
portpicker == 1.6.0
scipy == 1.11.3
requests == 2.31.0
packaging==23.2
setuptools==68.2.2

        --lockfile=requirements.txt:./requirements_lock_3_12.txt
        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.requirements.txt
        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/jax.requirements.txt

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

        "//tensorflow/compiler/mlir/lite/python:flatbuffer_to_mlir",
        "//tensorflow/compiler/mlir/lite/python:graphdef_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/python:jax_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/python:saved_model_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/quantization/lite:quantize_model",