create artifact list dynamically

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -62,7 +62,11 @@ jobs:
     strategy:
       matrix:
         runner: ["windows-x86-n2-64", "linux-x86-n2-16", "linux-arm64-t2a-16"]
-        artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
+        artifact: >-
+                  ${{ fromJSON((inputs.build_jax && '["jax"]') || '[]') }} +
+                  ${{ fromJSON((inputs.build_jaxlib && '["jaxlib"]') || '[]') }} +
+                  ${{ fromJSON((inputs.build_jax_cuda_pjrt && '["jax-cuda-pjrt"]') || '[]') }} +
+                  ${{ fromJSON((inputs.build_jax_cuda_plugin && '["jax-cuda-plugin"]') || '[]')  }}
         python: ["3.10"] #, "3.11", "3.12"]
         # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each
         # Python version.
@@ -97,46 +101,20 @@ jobs:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
       JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}"
 
-    if: inputs.build_jax || inputs.build_jaxlib || inputs.build_jax_cuda_plugin || inputs.build_jax_cuda_pjrt
     steps:
       - uses: actions/checkout@v3
-        if: >-
-            (inputs.build_jax && matrix.artifact == 'jax') ||
-            (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
-            (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
-            (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-        if: >-
-          (inputs.build_jax && matrix.artifact == 'jax') ||
-          (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
-          (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
-          (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
       - name: Build ${{ matrix.artifact }}
-        if: >-
-            (inputs.build_jax && matrix.artifact == 'jax') ||
-            (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
-            (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
-            (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
         run: ./ci/build_artifacts.sh "${{ matrix.artifact }}"
       - name: Set Platform
-        if: >-
-            (inputs.build_jax && matrix.artifact == 'jax') ||
-            (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
-            (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
-            (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
         run: |
           echo "PLATFORM=$(uname)_$(uname -m)" >> $GITHUB_ENV
       - name: Upload artifacts to GCS bucket
         # Upload if requested and one of the artifacts was built
-        if: >-
-          inputs.upload_artifacts &&
-          (inputs.build_jax && matrix.artifact == 'jax') ||
-          (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
-          (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
-          (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
+        if: inputs.upload_artifacts
         run:  ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r $(pwd)/dist gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination }}"/$PLATFORM