Avoid depending on JAX internals

.github/workflows/pytest_and_autopublish.yml

@@ -14,7 +14,7 @@ jobs:
     # Install deps
     - uses: actions/setup-python@v4
       with:
-        python-version: '3.10'
+        python-version: '3.11'
 
     - run: sudo apt install -y pandoc gsfonts
     - run: pip --version
@@ -41,14 +41,13 @@ jobs:
     # Install deps
     - uses: actions/setup-python@v4
       with:
-        python-version: '3.10'
+        python-version: '3.11'
         # Uncomment to cache of pip dependencies (if tests too slow)
         # cache: pip
         # cache-dependency-path: '**/pyproject.toml'
 
-    - run:
-        pip install -e .[dev]
-        pip install "git+https://github.com/jax-ml/oryx.git@b59ab020780cd53d488bc7dcad3696be9fdca0a5"
+    - run: pip install -e .[dev]
+    - run: pip install "git+https://github.com/jax-ml/oryx.git"
 
     # Run tests (in parallel)
     - name: Run core tests

.pylintrc

@@ -137,6 +137,7 @@ disable=abstract-method,
         too-many-instance-attributes,
         too-many-locals,
         too-many-nested-blocks,
+        too-many-positional-arguments,
         too-many-public-methods,
         too-many-return-statements,
         too-many-statements,

coix/oryx.py

@@ -250,7 +250,6 @@ def distribution_rule(state, *args, **kwargs):
     _, dist = jax.tree.unflatten(kwargs["in_tree"], flat_args)
     dist_flat, dist_tree = jax.tree.flatten(dist)
     state[name] = {dist_tree: dist_flat}
-  args = jax.tree.map(jax.core.raise_as_much_as_possible, args)
   return random_variable_p.bind(*args, **kwargs), state
 
 

coix/oryx_test.py

@@ -113,7 +113,7 @@ def expected_fn(x):
 
 def test_observed():
   def model(a):
-    return coryx.rv(dist.Delta(2.0, 3.0), obs=1.0, name="x") + a
+    return coryx.rv(dist.Delta(a, 3.0), obs=1.0, name="x") + a
 
   _, trace, _ = coix.traced_evaluate(model)(2.0)
   assert "x" in trace

docs/conf.py

@@ -234,7 +234,6 @@
 # a list of builtin themes.
 #
 html_theme = "sphinx_rtd_theme"
-html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the