raise TypeError if device_buffer is a tuple

jakevdp · Sep 22, 2020 · 699e46a · 699e46a
1 parent e3cfaf6
commit 699e46a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
@@ -969,6 +969,9 @@ class DeviceArray:
   def __init__(self, aval: core.ShapedArray, device: Optional[Device],
                lazy_expr: lazy.LazyExpr,
                device_buffer: PyLocalBuffer):
+    # TODO(jakevdp): remove this check once multi-buf change is baked.
+    if isinstance(device_buffer, Sequence):
+      raise TypeError(f"device_buffer should be a Buffer instance; got {device_buffer}")
     self.aval = aval
     self.device_buffer = device_buffer
     self._device = device