Fix provenance logic for jax 0.2.28 (pyro-ppl#1320)

themrzmaster · Feb 4, 2022 · 9719466 · 9719466
1 parent a69344f
commit 9719466
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
@@ -26,6 +26,9 @@ def process_primitive(self, primitive, tracers, params):
             out_tracers = out_tracers if primitive.multiple_results else [out_tracers]
             for t in out_tracers:
                 t.aval.named_shape["_provenance"] = out_provenance
+                # Also update provenance of the cached tracer -> aval dict.
+                aval_cache = self.frame.tracer_to_var[id(t)].aval
+                aval_cache.named_shape["_provenance"] = out_provenance
             out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]
         return out_tracers
 

diff --git a/setup.py b/setup.py
@@ -9,7 +9,7 @@
 from setuptools import find_packages, setup
 
 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
-_jax_version_constraints = ">=0.2.13,<0.2.28"
+_jax_version_constraints = ">=0.2.13"
 _jaxlib_version_constraints = ">=0.1.65"
 
 # Find version