Update JAX and TFLite model artifacts

iree-org · Mar 12, 2024 · 6026c3d · 6026c3d
1 parent 120cf01
commit 6026c3d
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 15 deletions.
diff --git a/.github/workflows/run_comparative_benchmark.yml b/.github/workflows/run_comparative_benchmark.yml
@@ -242,18 +242,19 @@ jobs:
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
           gcloud storage cp "${RESULTS_PATH}" "${RESULTS_GCS_DIR}/"
-      - name: "Benchmarking JAX-IREE:CPU"
-        env:
-          JAX_IREE_RESULTS_JSON: jax-iree.json
-          RESULTS_GCS_DIR: ${{ steps.setup.outputs.results-gcs-dir }}
-        run: |
-          RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
-          docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:b9b98da7bcc5e431800ff798a6dcc394b1838a9ed3d695f5cd0dac3510fc8c8d" \
-            ./comparative_benchmark/jax/benchmark_iree.sh \
-              "${TARGET_DEVICE}"\
-              "${RESULTS_PATH}"
-          gcloud storage cp "${RESULTS_PATH}" "${RESULTS_GCS_DIR}/"
+#      # Disabled due to https://github.com/openxla/openxla-benchmark/issues/169.
+#      - name: "Benchmarking JAX-IREE:CPU"
+#        env:
+#          JAX_IREE_RESULTS_JSON: jax-iree.json
+#          RESULTS_GCS_DIR: ${{ steps.setup.outputs.results-gcs-dir }}
+#        run: |
+#          RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
+#          docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
+#            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:b9b98da7bcc5e431800ff798a6dcc394b1838a9ed3d695f5cd0dac3510fc8c8d" \
+#            ./comparative_benchmark/jax/benchmark_iree.sh \
+#              "${TARGET_DEVICE}"\
+#              "${RESULTS_PATH}"
+#          gcloud storage cp "${RESULTS_PATH}" "${RESULTS_GCS_DIR}/"
 #      # Disabled in favor of JAX and reducing maintenance burden.
 #      - name: "Benchmarking TF-XLA:CPU"
 #        env:

diff --git a/common_benchmark_suite/openxla/benchmark/comparative_suite/jax/model_definitions.py b/common_benchmark_suite/openxla/benchmark/comparative_suite/jax/model_definitions.py
@@ -10,7 +10,7 @@
 from openxla.benchmark import def_types
 from openxla.benchmark.comparative_suite import utils
 
-PARENT_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/jax/jax_models_0.4.23_1705868085"
+PARENT_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/jax/jax_models_0.4.25_1709787220"
 ARTIFACTS_DIR_URL_TEMPLATE = string.Template(PARENT_GCS_DIR + "/${name}")
 
 T5_JAX_IMPL = def_types.ModelImplementation(

diff --git a/common_benchmark_suite/openxla/benchmark/comparative_suite/tflite/model_definitions.py b/common_benchmark_suite/openxla/benchmark/comparative_suite/tflite/model_definitions.py
@@ -10,7 +10,7 @@
 from openxla.benchmark import def_types
 from openxla.benchmark.comparative_suite import utils
 
-PARENT_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/tflite/tflite_models_1706739936"
+PARENT_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/tflite/tflite_models_1710104772"
 ARTIFACTS_DIR_URL_TEMPLATE = string.Template(PARENT_GCS_DIR + "/${name}")
 
 TFLITE_MODEL_IMPL = def_types.ModelImplementation(
@@ -20,7 +20,7 @@
     module_path=f"{utils.MODELS_MODULE_PATH}.tflite.tflite_model",
 )
 
-JAX_MODELS_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/jax/jax_models_0.4.23_1706594181"
+JAX_MODELS_GCS_DIR = "https://storage.googleapis.com/iree-model-artifacts/jax/jax_models_0.4.25_1709787220"
 
 BERT_BASE_FP32_TFLITE_I32_INPUT_SEQUENCE_TEMPLATE = utils.ModelTemplate(
     name=utils.SEQ_LEN_NAME("BERT_BASE_FP32_TFLITE_I32"),