fix

.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 }}!"
+          echo "Reason is ${{ github.event.inputs.reason }}!"

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 .`
 

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:]