Enable colored output in pytest under github actions.

.github/workflows/ci-build.yaml

@@ -112,6 +112,7 @@ jobs:
         JAX_THREEFRY_PARTITIONABLE: ${{ matrix.prng-upgrade }}
         JAX_ENABLE_CHECKS: true
         JAX_SKIP_SLOW_TESTS: true
+        PY_COLORS: 1
       run: |
         pip install -e .
         echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
@@ -158,6 +159,7 @@ jobs:
       env:
         XLA_FLAGS: "--xla_force_host_platform_device_count=8"
         JAX_ARRAY: 1
+        PY_COLORS: 1
       run: |
         pytest -n auto --tb=short docs
         pytest -n auto --tb=short --doctest-modules jax --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/interpreters/mlir.py --ignore=jax/_src/iree.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -55,6 +55,7 @@ jobs:
         env:
           JAX_PLATFORMS: tpu,cpu
           JAX_USE_PJRT_C_API_ON_TPU: ${{ matrix.pjrt }}
+          PY_COLORS: 1
         run: |
           # Run single-accelerator tests in parallel
           JAX_ENABLE_TPU_XDIST=true python3 -m pytest -n=4 --tb=short \

.github/workflows/upstream-nightly.yaml

@@ -65,6 +65,7 @@ jobs:
           JAX_ENABLE_X64: true
           JAX_ENABLE_CHECKS: true
           JAX_SKIP_SLOW_TESTS: true
+          PY_COLORS: 1
         run: |
           echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
           echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"

.github/workflows/wheel_win_x64.yml

@@ -50,6 +50,7 @@ jobs:
         env:
           JAX_ENABLE_CHECKS: true
           JAX_SKIP_SLOW_TESTS: true
+          PY_COLORS: 1
         run: |
           python -m pip install -e ${{ github.workspace }}
           python -m pip install --no-index --find-links ${{ github.workspace }}\dist jaxlib

.github/workflows/windows_ci.yml

@@ -59,6 +59,7 @@ jobs:
         env:
           JAX_ENABLE_CHECKS: true
           JAX_SKIP_SLOW_TESTS: true
+          PY_COLORS: 1
         run: |
           cd jax
           python -m pip install -e ${{ github.workspace }}\jax