Skip to content

JAX FP8 matmul fusion tutorial. #399

JAX FP8 matmul fusion tutorial.

JAX FP8 matmul fusion tutorial. #399

Workflow file for this run

name: tests
env:
GIT_MAIN_BRANCH: "main"
# Controls when the workflow will run.
on:
push:
branches: [ "main" ]
pull_request:
branches: [ "main" ]
# Allows you to run this workflow manually from the Actions tab.
workflow_dispatch:
jobs:
lint_and_typecheck:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.11.0
with:
access_token: ${{ github.token }}
if: ${{github.ref != 'refs/head/main'}}
- uses: actions/checkout@v3
- name: Set up Python 3.10
uses: actions/setup-python@v4
with:
python-version: "3.10"
- uses: pre-commit/action@v3.0.0
unit_tests:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.11.0
with:
access_token: ${{ github.token }}
if: ${{github.ref != 'refs/head/main'}}
- uses: actions/checkout@v3
- name: Update pip
id: pip-cache
run: |
python3 -m pip install --upgrade pip
- name: Local install & test requirements
run: |
pip3 install -e ./
pip3 install -r ./test-requirements.txt
# Run repository unit tests on latest JAX
- name: Run unit tests JAX latest
run: |
pytest --tb=short -v --log-cli-level=INFO ./
- name: JAX 0.3.16 installation
run: |
pip3 install numpy==1.24.3 scipy==1.10.1
pip3 install chex==0.1.6 jax==0.3.16 jaxlib==0.3.15 -f https://storage.googleapis.com/jax-releases/jax_releases.html
# Run repository unit tests on JAX 0.3
- name: Run unit tests JAX 0.3.16
run: |
pytest --tb=short -v --log-cli-level=INFO ./