Bfloat16 jnp.allclose raises Illegal Instruction Exception

[aturker-synnada]

Description

Hello, I started using bfloat16 in my JAX project, but the GitHub CI workflow began raising an Illegal instruction exception. Locally, I am using macOS 15 with an M2 Pro chip, and everything works perfectly. However, in the GitHub CI environment, the error occurs. I suspect the issue might be related to jnp.allclose. Interestingly, if I remove the line jax.config.update("jax_enable_x64", True), everything works correctly.

The script:

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)


def example_fn():
    x = jnp.array([1, 2, 3], dtype=jnp.bfloat16)
    y = jnp.array([1.0, 2, 3], dtype=jnp.bfloat16)
    jnp.allclose(x, y, 1e-5, 1e-5)

    print("example_fn done")


if __name__ == "__main__":
    example_fn()

The workflow:

name: CI-MacOS
on:
  push:
    branches: [ dev, main ]
  pull_request:
    branches: [ dev, main ]
env:
  CI: true
jobs:
  build:
    runs-on: macos-14
    strategy:
      matrix:
        python-version: [3.12]
    steps:
      - name: Check out repository
        uses: actions/checkout@v2
      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v2
        with:
          python-version: ${{ matrix.python-version }}
      - name: Install Python dependencies
        run: |
          python3 -m pip install --upgrade pip
          pip install --upgrade "jax[cpu]"
      - name: Execute script
        run: python3 script.py
      - name: Upload results to Codecov
        uses: codecov/codecov-action@v4
        with:
          token: ${{ secrets.CODECOV_TOKEN }}      

Workflow output:

Run python3 script.py
  python3 script.py
  shell: /bin/bash -e {0}
  env:
    CI: true
    pythonLocation: /Users/runner/hostedtoolcache/Python/3.12.8/arm64
/Users/runner/work/_temp/ffb44547-aa7b-4e70-8faa-f264c0fdb4f1.sh: line 1:  1886 Illegal instruction: 4  python3 script.py

System info (python version, jaxlib version, accelerator, etc.)

jaxlib: 0.4.38
numpy: 2.2.1
python: 3.12.8 (v3.12.8:2dc476bcb91, Dec 3 2024, 14:43:19) [Clang 13.0.0 (clang-1300.0.29.30)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='Mac-1736162455660.local', release='23.6.0', version='Darwin Kernel Version 23.6.0: Fri Nov 15 15:13:40 PST 2024; root:xnu-10063.141.1.702.7~1/RELEASE_ARM64_VMAPPLE', machine='arm64')

[jakevdp]

Can you paste the traceback here? (I appreciate the repro, but it doesn't repro locally for me and it would help to see the error text without having to go through the process of setting up a new github action!)

[aturker-synnada]

Thanks for the quick response. I cannot reproduce the error locally; I believe it is related to something with the virtual M1 processor and JAX-XLA. The Python does not provide a traceback in the error (I am not sure why). The example workflow run can be seen here. This workflow was worked with above code(did not import the library).

[aturker-synnada]

I created an empty repository to reproduce this issue. See the repo. @jakevdp