Also do not patch MaskedAutoregressiveFlow.

PiperOrigin-RevId: 599551554
jax-ml · Jan 18, 2024 · fd6c9cf · fd6c9cf
1 parent ebcdedf
commit fd6c9cf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/oryx/core/interpreters/inverse/rules.py b/oryx/core/interpreters/inverse/rules.py
@@ -40,7 +40,7 @@
 for name in tfb.__all__:
   bij = getattr(tfb, name)
   if inspect.isclass(bij) and issubclass(bij, tfb.Bijector):
-    if bij not in {tfb.Bijector, tfb.Composition}:
+    if bij not in {tfb.Bijector, tfb.Composition, tfb.MaskedAutoregressiveFlow}:
       bijector_extensions.patch_bijector(bij)
 
 ildj_registry = inverse_core.ildj_registry