update jax_ort env

Signed-off-by: John <[email protected]>
onnx · Jul 27, 2023 · ef575b6 · ef575b6
1 parent 0b472a7
commit ef575b6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git 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
@@ -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