Skip to content

Replace deprecated jax.tree_* functions with jax.tree.* #3698

Replace deprecated jax.tree_* functions with jax.tree.*

Replace deprecated jax.tree_* functions with jax.tree.* #3698

Workflow file for this run

name: GitHub Action pytest
on:
workflow_dispatch:
push:
branches:
- main
pull_request:
branches:
- main
jobs:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Test Import
run: |
set -v
set -e
python -c "import uncertainty_baselines as ub"
python -c "import uncertainty_baselines as ub; print(ub.datasets)"
python -c "import uncertainty_baselines as ub; print(ub.models)"
- name: Run Pytest
# TODO(dusenberrymw,trandustin): Reenable bert tests.
run: |
set -v
set -e
python -m unittest -v uncertainty_baselines/datasets/random_test.py
python -m unittest -v $(find . -name 'uncertainty_baselines/models/*_test.py' -a ! -name 'uncertainty_baselines/models/bert*_test.py')
python -m unittest -v uncertainty_baselines/*_test.py
# NOTE: These JFT directory tests are split out to lower the total runtime.
test_jft_active_learning:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v active_learning_test.py
python -m unittest -v al_utils_test.py
test_jft_batchensemble:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v batchensemble_test.py
test_jft_deep_ensemble:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v deep_ensemble_test.py
test_jft_deterministic:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v deterministic_test.py
test_jft_heteroscedastic:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v heteroscedastic_test.py
test_jft_sngp:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v sngp_test.py
test_jft_utils:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- name: Install Python 3
uses: actions/setup-python@v1
with:
python-version: 3.7
- name: Pip Install
run: |
set -v
set -e
python -m pip install --upgrade pip
pip install -q -e .[tensorflow,jax,models,tests]
- name: Run Pytest
run: |
set -v
set -e
cd baselines/jft
python -m unittest -v checkpoint_utils_test.py
python -m unittest -v input_utils_test.py
python -m unittest -v ood_utils_test.py
python -m unittest -v preprocess_utils_test.py
python -m unittest -v train_utils_test.py