From fb78efba006f9206e8231a3948d96de1a9af7314 Mon Sep 17 00:00:00 2001
From: Zheng Zhao <zgbkdlm@gmail.com>
Date: Wed, 22 May 2024 13:53:36 +0200
Subject: [PATCH] fix

---
 .github/workflows/unittest.yml | 8 ++++----
 README.md                      | 2 +-
 tests/test_gibbs.py            | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index 60319c68..030d661e 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -40,15 +40,15 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install wheel
+          pip install jaxlib jax
           pip install -r requirements.txt
-          pip install jax[cpu]
           pip install torch --index-url https://download.pytorch.org/whl/cpu
-          python setup.py install
+          pip install -e .
       - name: Do the unit tests
         run: |
           pip install pytest
-          python -m pytest tests/
+          pytest tests/
       - name: Logging
         run: |
           echo "This is ${{ github.event.inputs.name }}!"
-          echo "Reason is ${{ github.event.inputs.reason }}!"
\ No newline at end of file
+          echo "Reason is ${{ github.event.inputs.reason }}!"
diff --git a/README.md b/README.md
index 756b5db6..539fbe96 100644
--- a/README.md
+++ b/README.md
@@ -21,7 +21,7 @@ inherent statistical biases (plural!).
 1. `git clone git@github.com:zgbkdlm/fbs.git`
 2. `cd fbs`
 3. `python venv ./venv && source venv/bin/activate`
-4. Install JAX in either GPU/CPU environment according to this official guidance https://placeholder.
+4. Install JAX in either GPU/CPU environment according to this official guidance https://github.com/google/jax?tab=readme-ov-file#installation.
 5. `pip install -r requirements.txt`
 6. `pip install -e .`
 
diff --git a/tests/test_gibbs.py b/tests/test_gibbs.py
index ce924e8c..aa457c07 100644
--- a/tests/test_gibbs.py
+++ b/tests/test_gibbs.py
@@ -115,7 +115,7 @@ def gibbs_kernel(key_, x0_, y0_, us_star_, bs_star_):
     for i in range(nsamples):
         key, subkey = jax.random.split(key)
         x0, us_star, bs_stars, acc = gibbs_kernel(subkey, x0, y0, us_star, bs_stars)
-        x0s[i] = x0
+        x0s[i] = jnp.squeeze(x0)
 
     x0s = x0s[burnin:]