[Oryx] Fix use of jax.core.literalable_types

PiperOrigin-RevId: 334964477
gregreen · Oct 2, 2020 · 900c08d · 900c08d
1 parent aaed4d3
commit 900c08d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/spinoffs/oryx/oryx/core/interpreters/unzip.py b/spinoffs/oryx/oryx/core/interpreters/unzip.py
@@ -179,8 +179,7 @@ def instantiate_const(self, tracer):
     if isinstance(pv, jax_core.AbstractValue):
       return tracer
     elif not pv:
-      if (isinstance(const, tuple(jax_core.literalable_types)) and
-          not onp.shape(const)):
+      if type(const) in jax_core.literalable_types and not onp.shape(const):  # pylint: disable=unidiomatic-typecheck
         return self.new_instantiated_literal(const)
       else:
         return self.new_instantiated_const(const)