From 0b472a720b6b79c8313d08f2c9cea9273960a8f6 Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Thu, 27 Jul 2023 21:43:42 +0000 Subject: [PATCH 1/6] update the jax jaxort env for op test. Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/development/Dockerfile | 6 ++++++ runtimes/jaxonnxruntime/stable/Dockerfile | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/runtimes/jaxonnxruntime/development/Dockerfile b/runtimes/jaxonnxruntime/development/Dockerfile index 5b9196f493..c47e21c146 100644 --- a/runtimes/jaxonnxruntime/development/Dockerfile +++ b/runtimes/jaxonnxruntime/development/Dockerfile @@ -29,5 +29,11 @@ RUN pip3 install git+https://github.com/google/jaxonnxruntime.git@main RUN pip3 install --upgrade pip setuptools wheel pytest +############## JAX and JAXONNXRuntime configs ########### +# ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py +ENV JAX_DEFAULT_DTYPE_BITS=64 +ENV JAX_NUMPY_RANK_PROMOTION=warn +ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 + CMD . /root/setup/docker-setup.sh && \ pytest /root/test/test_backend.py --onnx_backend=${ONNX_BACKEND} -k 'not _cuda' -v diff --git a/runtimes/jaxonnxruntime/stable/Dockerfile b/runtimes/jaxonnxruntime/stable/Dockerfile index ae48f8b855..6dc820b0b8 100644 --- a/runtimes/jaxonnxruntime/stable/Dockerfile +++ b/runtimes/jaxonnxruntime/stable/Dockerfile @@ -29,5 +29,11 @@ RUN pip3 install --no-cache-dir \ #################################################### RUN pip3 install --upgrade pip setuptools wheel pytest +############## JAX and JAXONNXRuntime configs ########### +# ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py +ENV JAX_DEFAULT_DTYPE_BITS=64 +ENV JAX_NUMPY_RANK_PROMOTION=warn +ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 + CMD . /root/setup/docker-setup.sh && \ pytest /root/test/test_backend.py --onnx_backend=${ONNX_BACKEND} -k 'not _cuda' -v From ef575b6377dfa997f4047e6741640322545ce897 Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Thu, 27 Jul 2023 21:58:45 +0000 Subject: [PATCH 2/6] update jax_ort env Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/development/Dockerfile | 1 + runtimes/jaxonnxruntime/stable/Dockerfile | 1 + 2 files changed, 2 insertions(+) diff --git a/runtimes/jaxonnxruntime/development/Dockerfile b/runtimes/jaxonnxruntime/development/Dockerfile index c47e21c146..0a4414d017 100644 --- a/runtimes/jaxonnxruntime/development/Dockerfile +++ b/runtimes/jaxonnxruntime/development/Dockerfile @@ -34,6 +34,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ENV JAX_DEFAULT_DTYPE_BITS=64 ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 +ENV JAXORT_NONZERO_USE_FULLY_PADDING=1 CMD . /root/setup/docker-setup.sh && \ pytest /root/test/test_backend.py --onnx_backend=${ONNX_BACKEND} -k 'not _cuda' -v diff --git a/runtimes/jaxonnxruntime/stable/Dockerfile b/runtimes/jaxonnxruntime/stable/Dockerfile index 6dc820b0b8..855cde0334 100644 --- a/runtimes/jaxonnxruntime/stable/Dockerfile +++ b/runtimes/jaxonnxruntime/stable/Dockerfile @@ -34,6 +34,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ENV JAX_DEFAULT_DTYPE_BITS=64 ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 +ENV JAXORT_NONZERO_USE_FULLY_PADDING=1 CMD . /root/setup/docker-setup.sh && \ pytest /root/test/test_backend.py --onnx_backend=${ONNX_BACKEND} -k 'not _cuda' -v From 3c5c0234b213d46fc2179a0a8711d832d9f3a7b7 Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Mon, 31 Jul 2023 13:47:49 -0700 Subject: [PATCH 3/6] Update runtimes/jaxonnxruntime/development/Dockerfile Co-authored-by: Jinning Li Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/development/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtimes/jaxonnxruntime/development/Dockerfile b/runtimes/jaxonnxruntime/development/Dockerfile index 0a4414d017..08e20e5127 100644 --- a/runtimes/jaxonnxruntime/development/Dockerfile +++ b/runtimes/jaxonnxruntime/development/Dockerfile @@ -31,7 +31,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ############## JAX and JAXONNXRuntime configs ########### # ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py -ENV JAX_DEFAULT_DTYPE_BITS=64 +ENV JAX_ENABLE_X64=True ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 ENV JAXORT_NONZERO_USE_FULLY_PADDING=1 From 2b7fb25fcdf726560a6426ec037a5843400f78b0 Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Mon, 31 Jul 2023 13:50:26 -0700 Subject: [PATCH 4/6] Update Dockerfile ENV JAX_DEFAULT_DTYPE_BITS=64 Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/development/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/runtimes/jaxonnxruntime/development/Dockerfile b/runtimes/jaxonnxruntime/development/Dockerfile index 08e20e5127..ab6e5831d3 100644 --- a/runtimes/jaxonnxruntime/development/Dockerfile +++ b/runtimes/jaxonnxruntime/development/Dockerfile @@ -31,6 +31,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ############## JAX and JAXONNXRuntime configs ########### # ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py +ENV JAX_DEFAULT_DTYPE_BITS=64 ENV JAX_ENABLE_X64=True ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 From 315b9f51a7698f93b10de881328075799f20d844 Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Mon, 31 Jul 2023 13:51:12 -0700 Subject: [PATCH 5/6] Apply suggestions from code review Co-authored-by: Jinning Li Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/stable/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtimes/jaxonnxruntime/stable/Dockerfile b/runtimes/jaxonnxruntime/stable/Dockerfile index 855cde0334..38a26a0d60 100644 --- a/runtimes/jaxonnxruntime/stable/Dockerfile +++ b/runtimes/jaxonnxruntime/stable/Dockerfile @@ -31,7 +31,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ############## JAX and JAXONNXRuntime configs ########### # ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py -ENV JAX_DEFAULT_DTYPE_BITS=64 +ENV JAX_ENABLE_X64=True ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0 ENV JAXORT_NONZERO_USE_FULLY_PADDING=1 From 2dec154036b5c4136483e95ea92b812477566e0b Mon Sep 17 00:00:00 2001 From: John <1200986+maxwillzq@users.noreply.github.com> Date: Mon, 31 Jul 2023 13:51:35 -0700 Subject: [PATCH 6/6] Update Dockerfile ENV JAX_DEFAULT_DTYPE_BITS=64 Signed-off-by: John <1200986+maxwillzq@users.noreply.github.com> --- runtimes/jaxonnxruntime/stable/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/runtimes/jaxonnxruntime/stable/Dockerfile b/runtimes/jaxonnxruntime/stable/Dockerfile index 38a26a0d60..549693242d 100644 --- a/runtimes/jaxonnxruntime/stable/Dockerfile +++ b/runtimes/jaxonnxruntime/stable/Dockerfile @@ -31,6 +31,7 @@ RUN pip3 install --upgrade pip setuptools wheel pytest ############## JAX and JAXONNXRuntime configs ########### # ref to https://github.com/google/jaxonnxruntime/blob/main/tests/onnx_ops_test.py +ENV JAX_DEFAULT_DTYPE_BITS=64 ENV JAX_ENABLE_X64=True ENV JAX_NUMPY_RANK_PROMOTION=warn ENV JAXORT_ONLY_ALLOW_INITIALIZERS_AS_STATIC_ARGS=0