fix typo in pxla

google · Feb 23, 2019 · 8a2fd90 · 8a2fd90
1 parent b1686a3
commit 8a2fd90
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
@@ -150,7 +150,7 @@ def axis_env_extend(name, size):
 def xla_split(c, axis_sizes, x):
   def _xla_split(shape, x):
     if shape.is_tuple():
-      elts = map(_xla_split, shape.tuple_shapes(), _xla_destructure(c, x))
+      elts = map(_xla_split, shape.tuple_shapes(), xla_destructure(c, x))
       return c.Tuple(*elts)
     else:
       size = onp.array(prod(axis_sizes), onp.uint32)
@@ -167,7 +167,7 @@ def _xla_split(shape, x):
 def xla_join(c, device_groups, x):
   def _xla_join(shape, x):
     if shape.is_tuple():
-      elts = map(_xla_join, shape.tuple_shapes(), _xla_destructure(c, x))
+      elts = map(_xla_join, shape.tuple_shapes(), xla_destructure(c, x))
       return c.Tuple(*elts)
     else:
       group_size = len(device_groups[0])