Delete references to removed jax.core attributes.

PiperOrigin-RevId: 451174810
google-deepmind · May 26, 2022 · 883be30 · 883be30
1 parent f2ef70c
commit 883be30
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/distrax/_src/utils/transformations.py b/distrax/_src/utils/transformations.py
@@ -247,7 +247,6 @@ def read(var):
   def write(var, val):
     env[var] = val
 
-  write(jax.core.unitvar, jax.core.unit)
   jax.api_util.safe_map(write, jaxpr.outvars, args)
   jax.api_util.safe_map(write, jaxpr.constvars, consts)