xnd 0.2.0dev6 → 0.2.0dev7

data/ext/ruby_xnd/xnd/python/xnd/__init__.py

@@ -43,14 +43,14 @@ data.
 Importing PEP-3118 buffers is supported.
 """
 
+from ._version import __version__
 
 # Ensure that libndtypes is loaded and initialized.
 from ndtypes import ndt, instantiate, MAX_DIM
-from ._xnd import Xnd, XndEllipsis
-from itertools import accumulate
+from ._xnd import Xnd, XndEllipsis, data_shapes, _typeof
 from .contrib.pretty import pretty
 
-__all__ = ['xnd', 'XndEllipsis', 'typeof', '_typeof']
+__all__ = ['xnd', 'array', 'XndEllipsis', 'typeof']
 
 
 # ======================================================================
@@ -99,7 +99,7 @@ class xnd(Xnd):
     """
 
     def __new__(cls, value, *, type=None, dtype=None, levels=None,
-                typedef=None, dtypedef=None):
+                typedef=None, dtypedef=None, device=None):
         if (type, dtype, levels, typedef, dtypedef).count(None) < 2:
             raise TypeError(
                 "the 'type', 'dtype', 'levels' and 'typedef' arguments are "
@@ -124,7 +124,13 @@ class xnd(Xnd):
             type = typeof(value, dtype=dtype)
         else:
             type = typeof(value)
-        return super().__new__(cls, type=type, value=value)
+
+        if device is not None:
+            name, no = device.split(":")
+            no = -1 if no == "managed" else int(no)
+            device = (name, no)
+
+        return super().__new__(cls, type=type, value=value, device=device)
 
     def __repr__(self):
         value = self.short_value(maxshape=10)
@@ -134,157 +140,636 @@ class xnd(Xnd):
         fmt = fmt.replace("\n", "\n   ")
         return "xnd%s" % fmt
 
+    def __reduce__(self):
+        b =  self.serialize()
+        return (xnd.deserialize, (b,))
+
+    def copy_contiguous(self, dtype=None):
+        if isinstance(dtype, str):
+            dtype = ndt(dtype)
+        return super().copy_contiguous(dtype=dtype)
+
+    def reshape(self, *args, order=None):
+        return super()._reshape(args, order=order)
+
+    def serialize(self):
+        if not self.type.is_c_contiguous() and \
+           not self.type.is_f_contiguous() and \
+           not self.type.is_var_contiguous():
+            self = self.copy_contiguous()
+        return self._serialize()
+
     @classmethod
-    def unsafe_from_data(cls, obj=None, type=None):
+    def empty(cls, type=None, device=None):
+        if device is not None:
+            name, no = device.split(":")
+            no = -1 if no == "managed" else no
+            device = (name, int(no))
+
+        return super(xnd, cls).empty(type, device)
+
+    @classmethod
+    def from_buffer_and_type(cls, obj=None, type=None):
         """Return an xnd object that obtains memory from 'obj' via the
-           buffer protocol.  The buffer protocol's type is overridden by
-          'type'.  No safety checks are performed, the user is responsible
-           for passing a suitable type.
+           buffer protocol.  'obj' must be a simple writable buffer with
+           format 'B'.  The xnd object uses the provided type, which must
+           have the same data size as 'obj'.
         """
         if isinstance(type, str):
             type = ndt(type)
-        return cls._unsafe_from_data(obj, type)
+        return super().from_buffer_and_type(obj, type)
+
+def typeof(v, dtype=None):
+    if isinstance(dtype, str):
+        dtype = ndt(dtype)
+    return _typeof(v, dtype=dtype, shortcut=True)
 
 
 # ======================================================================
-#                            Type inference
+#                              array object
 # ======================================================================
 
-def typeof(value, *, dtype=None):
-    return ndt(_typeof(value, dtype=dtype))
+def _convert_smallest(v, device=None):
+    """Inefficient hack to make dask work (this needs to be in _typeof)."""
+    try:
+        return array.from_buffer(v)
+    except (TypeError, BufferError):
+        pass
+    x = array(v, device=device)
+    if x.type.hidden_dtype == ndt("int64"):
+        for dtype in ("int8", "int16", "int32"):
+            try:
+                return array(v, dtype=dtype, device=device)
+            except:
+                continue
+    return x
+
+class array(xnd):
+    """Extended array type that relies on gumath for the array functions."""
+
+    _functions = None
+    _cuda = None
+    _np = None
+
+    @property
+    def shape(self):
+        return self.type.shape
+
+    @property
+    def strides(self):
+        return self.type.strides
+
+    @property
+    def T(self):
+        return self.transpose()
+
+    def tolist(self):
+        return self.value
+
+    def _get_module(self):
+        if self.device == "cuda:managed":
+            if array._cuda is None:
+                import gumath.cuda
+                array._cuda = gumath.cuda
+            return array._cuda
+        else:
+            if array._functions is None:
+                import gumath.functions
+                array._functions = gumath.functions
+            return array._functions
+
+    def _get_numpy(self):
+        if array._np is None:
+            import numpy
+            array._np = numpy
+        return array._np
+
+    def _convert(self, other, raiseit=False):
+        if isinstance(other, array):
+            if other.device != self.device:
+                raise NotImplementedError("arrays must be on the same device")
+            return other
+        try:
+            return _convert_smallest(other, device=self.device)
+        except TypeError:
+            if raiseit:
+                raise TypeError("unable to convert %s to array" % other)
+            return NotImplemented
+
+    def _call_unary(self, name, out=None):
+        m = self._get_module()
+        return getattr(m, name)(self, out=out, cls=array)
+
+    def _call_binary(self, name, other, out=None, raiseit=False):
+        other = self._convert(other, raiseit)
+        m = self._get_module()
+        return getattr(m, name)(self, other, out=out, cls=array)
+
+    def _call_binary_np(self, name, other, out=None, raiseit=False):
+        """redirect unimplemented binary methods"""
+        np = self._get_numpy()
+        other = self._convert(other, raiseit)
+        x = getattr(np, name)(self, other, out=out)
+        if out is not None:
+            return out
+        return array.from_buffer(x)
+
+    __array_priority__ = 1000
+
+    @property
+    def __array_interface__(self):
+        shape = self.shape
+        typestr = ndt.to_format(self.dtype)
+        return dict(shape=shape, typestr=typestr, version=3)
+
+    def __array__(self, dtype=None):
+        if dtype is not None:
+            np = self._get_numpy()
+            x = np.array([], dtype=dtype)
+            t = ndt.from_format(memoryview(x).format)
+            return self.copy(dtype=t)
+        return self
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        np = self._get_numpy()
+
+        def conv_args(t):
+            if isinstance(t, tuple):
+                return tuple(conv_args(v) for v in t)
+            if isinstance(t, list):
+                return list(conv_args(v) for v in t)
+            elif isinstance(t, array):
+                return np.array(t, copy=False)
+            elif isinstance(t, np.ndarray):
+                return t
+            else:
+                raise NotImplementedError
+
+        np_self = np.array(self, copy=False)
+        try:
+            np_inputs = conv_args(inputs)
+        except NotImplementedError:
+            return NotImplemented
+
+        np_kwargs = kwargs.copy()
+        out = np_kwargs.pop("out", None)
+        if out is not None:
+            try:
+                np_out = conv_args(out)
+            except NotImplementedError:
+                return NotImplemented
+            np_kwargs["out"] = np_out
+
+        np_res = np_self.__array_ufunc__(ufunc, method, *np_inputs, **np_kwargs)
+        if np_res is NotImplemented:
+            return NotImplemented
+
+        if out is None:
+            if isinstance(np_res, tuple):
+                out = tuple(array.from_buffer(v) for v in np_res)
+            elif isinstance(np_res, list):
+                out = list(array.from_buffer(v) for v in np_res)
+            else:
+                out = array.from_buffer(np_res)
 
-def _choose_dtype(lst):
-    for x in lst:
-        if x is not None:
-            return _typeof(x)
-    return "float64"
+        return out
 
-def _typeof(value, *, dtype=None):
-    """Infer the type of a Python value.  Only a subset of Datashape is
-       supported.  In general, types need to be explicitly specified when
-       creating xnd objects.
-    """
-    if isinstance(value, list):
-        data, shapes = data_shapes(value)
-        opt = None in data
+    def __array_function__(self, func, types, args, kwargs):
+        np = self._get_numpy()
 
-        if dtype is None:
-            if not data:
-                dtype = 'float64'
+        def conv_types(t):
+            if isinstance(t, tuple):
+                return tuple(conv_types(v) for v in t)
+            if isinstance(t, list):
+                return list(conv_types(v) for v in t)
+            elif issubclass(t, (array, np.ndarray)):
+                return np.ndarray
+            else:
+                raise NotImplementedError
+
+        def conv_args(t, allow_ndarray=True):
+            if isinstance(t, tuple):
+                return tuple(conv_args(v) for v in t)
+            if isinstance(t, list):
+                return list(conv_args(v) for v in t)
+            elif isinstance(t, array):
+                return np.array(t, copy=False)
+            elif isinstance(t, np.ndarray):
+                return t if allow_ndarray else NotImplemented
             else:
-                dtype = _choose_dtype(data)
-                for x in data:
-                    if x is not None:
-                        t = _typeof(x)
-                        if t != dtype:
-                            raise ValueError("dtype mismatch: have %s and %s" % (dtype, t))
+                return t
+
+        def conv_res(t):
+            if isinstance(t, tuple):
+                return tuple(conv_res(v) for v in t)
+            if isinstance(t, list):
+                return list(conv_res(v) for v in t)
+            elif isinstance(t, np.ndarray):
+                return array.from_buffer(t)
+            else:
+                try:
+                    return array.from_buffer(memoryview(t))
+                except TypeError:
+                    return t
+
+        np_self = np.array(self, copy=False)
+        try:
+            np_types = conv_types(types)
+        except NotImplementedError:
+            return NotImplemented
+
+        try:
+            np_args = conv_args(args)
+        except NotImplementedError:
+            return NotImplemented
+
+        np_kwargs = kwargs.copy()
+        out = np_kwargs.pop("out", None)
+        if out is not None:
+            try:
+                np_out = conv_args(out, allow_ndarray=False)
+            except NotImplementedError:
+                return NotImplemented
+            np_kwargs["out"] = np_out
+
+        np_res = np_self.__array_function__(func, np_types, np_args, np_kwargs)
+        if np_res is NotImplemented:
+            return NotImplemented
+
+        if out is None:
+            return conv_res(np_res)
+        else:
+            return out
 
-            if opt:
-                dtype = '?' + dtype
+    def __repr__(self):
+        value = self.short_value(maxshape=10)
+        fmt = pretty((value, "@type='%s'@" % self.type), max_width=120)
+        fmt = fmt.replace('"@', "")
+        fmt = fmt.replace('@"', "")
+        fmt = fmt.replace("\n", "\n     ")
+        return "array%s" % fmt
 
-        t = dtype
-        var = any(len(set(lst)) > 1 or None in lst for lst in shapes)
-        for lst in shapes:
-            opt = None in lst
-            lst = [0 if x is None else x for x in lst]
-            t = add_dim(opt=opt, shapes=lst, typ=t, use_var=var)
+    def __copy__(self):
+        return self.copy()
 
-        return t
+    def __deepcopy__(self, memo):
+        return self.copy()
 
-    elif dtype is not None:
-        raise TypeError("dtype argument is only supported for arrays")
+    def __bool__(self):
+        np = self._get_numpy()
+        return bool(np.array(self, copy=False))
 
-    elif isinstance(value, tuple):
-        return "(" + ", ".join([_typeof(x) for x in value]) + ")"
+    def __neg__(self):
+        return self._call_unary("negative")
 
-    elif isinstance(value, dict):
-        if all(isinstance(k, str) for k in value):
-            return "{" + ", ".join(["%s: %s" % (k, _typeof(v)) for k, v in value.items()]) + "}"
-        raise ValueError("all dict keys must be strings")
+    def __pos__(self):
+        return self.copy()
 
-    elif value is None:
-        return '?float64'
+    def __abs__(self):
+        return self._call_unary("abs")
 
-    elif isinstance(value, float):
-        return 'float64'
+    def __invert__(self):
+        return self._call_unary("invert")
 
-    elif isinstance(value, complex):
-        return 'complex128'
+    def __complex__(self):
+        np = self._get_numpy()
+        return complex(np.array(self, copy=False))
 
-    elif isinstance(value, int):
-        return 'int64'
+    def __int__(self):
+        np = self._get_numpy()
+        return int(np.array(self, copy=False))
 
-    elif isinstance(value, str):
-        return 'string'
+    def __oct__(self):
+        np = self._get_numpy()
+        return oct(np.array(self, copy=False))
 
-    elif isinstance(value, bytes):
-        return 'bytes'
+    def __hex__(self):
+        np = self._get_numpy()
+        return hex(np.array(self, copy=False))
 
-    else:
-        raise ValueError("cannot infer type for %r" % value)
+    def __float__(self):
+        np = self._get_numpy()
+        return float(np.array(self, copy=False))
 
+    def __index__(self):
+        np = self._get_numpy()
+        return np.array(self, copy=False).__index__()
 
-def add_dim(*, opt=False, shapes=None, typ=None, use_var=False):
-    """Construct a new dimension type based on the list of 'shapes' that
-       are present in a dimension.
-    """
-    if use_var:
-        offsets = [0] + list(accumulate(shapes))
-        return "%svar(offsets=%s) * %s" % ('?' if opt else '', offsets, typ)
-    else:
-        n = len(set(shapes))
-        assert n <= 1 and not None in shapes
-        shape = 0 if n == 0 else shapes[0]
-        return "%d * %s" % (shape, typ)
-
-def data_shapes(tree):
-    """Extract array data and dimension shapes from a nested list. The
-       list may contain None for missing data or dimensions.
-
-       >>> data_shapes([[0, 1], [2, 3, 4], [5, 6, 7, 8]])
-       ([0, 1, 2, 3, 4, 5, 6, 7, 8], [[2, 3, 4], [3]])
-                     ^                    ^       ^
-                     |                    |       `--- ndim=2: single shape 3
-                     |                    `-- ndim=1: shapes 2, 3, 4
-                     `--- ndim=0: extracted array data
-    """
-    acc = [[] for _ in range(MAX_DIM+1)]
-    min_level = MAX_DIM + 1
-    max_level = 0
-
-    def search(level, a):
-        nonlocal min_level, max_level
-
-        if level > MAX_DIM:
-            raise ValueError("too many dimensions")
-
-        current = acc[level]
-        if a is None:
-            current.append(a)
-        elif isinstance(a, list):
-            current.append(len(a))
-            next_level = level + 1
-            max_level = max(next_level, max_level)
-            if not a:
-                min_level = min(next_level, min_level)
-            else:
-                for item in a:
-                    search(level+1, item)
-        else:
-            acc[max_level].append(a)
-            min_level = min(level, min_level)
+    def __floor__(self):
+        np = self._get_numpy()
+        return np.array(self, copy=False).__floor__()
+
+    def __ceil__(self):
+        np = self._get_numpy()
+        return np.array(self, copy=False).__ceil__()
+
+    def __eq__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("equal", other)
+
+    def __ne__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("not_equal", other)
+
+    def __lt__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("less", other)
+
+    def __le__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("less_equal", other)
+
+    def __ge__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("greater_equal", other)
+
+    def __gt__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return self._call_binary("greater", other)
+
+    def __add__(self, other):
+        return self._call_binary("add", other)
+
+    def __radd__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("add", self)
+
+    def __sub__(self, other):
+        return self._call_binary("subtract", other)
+
+    def __rsub__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("subtract", self)
+
+    def __mul__(self, other):
+        return self._call_binary("multiply", other)
+
+    def __rmul__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("multiply", self)
+
+    def __matmul__(self, other):
+        return self._call_binary_np("matmul", other)
+
+    def __rmatmul__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary_np("matmul", self)
+
+    def __truediv__(self, other):
+        return self._call_binary("divide", other)
+
+    def __rtruediv__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("divide", self)
+
+    def __floordiv__(self, other):
+        return self._call_binary("floor_divide", other)
+
+    def __rfloordiv__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("floor_divide", self)
+
+    def __mod__(self, other):
+        return self._call_binary("remainder", other)
+
+    def __rmod__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("remainder", self)
+
+    def __divmod__(self, other):
+        return self._call_binary("divmod", other)
+
+    def __rdivmod__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("divmod", self)
+
+    def __pow__(self, other):
+        return self._call_binary("power", other)
+
+    def __rpow__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("power", self)
+
+    def __lshift__(self, other):
+        return self._call_binary_np("left_shift", other)
+
+    def __rlshift__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary_np("left_shift", self)
+
+    def __rshift__(self, other):
+        return self._call_binary_np("right_shift", other)
+
+    def __rrshift__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary_np("right_shift", self)
+
+    def __and__(self, other):
+        return self._call_binary("bitwise_and", other)
+
+    def __rand__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("bitwise_and", self)
+
+    def __or__(self, other):
+        return self._call_binary("bitwise_or", other)
+
+    def __ror__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("bitwise_or", self)
+
+    def __xor__(self, other):
+        return self._call_binary("bitwise_xor", other)
+
+    def __rxor__(self, other):
+        other = self._convert(other)
+        if other is NotImplemented:
+            return other
+
+        return other._call_binary("bitwise_xor", self)
+
+    def __iadd__(self, other):
+        return self._call_binary("add", other, out=self)
+
+    def __isub__(self, other):
+        return self._call_binary("subtract", other, out=self)
+
+    def __imul__(self, other):
+        return self._call_binary("multiply", other, out=self)
+
+    def __imatmul__(self, other):
+        return self._call_binary_np("matmul", other, out=self)
+
+    def __itruediv__(self, other):
+        return self._call_binary("divide", other, out=self)
+
+    def __ifloordiv__(self, other):
+        return self._call_binary("floor_divide", other, out=self)
+
+    def __imod__(self, other):
+        return self._call_binary("remainder", other, out=self)
+
+    def __ipow__(self, other):
+        return self._call_binary("power", other, out=self)
+
+    def __ilshift__(self, other):
+        return self._call_binary_np("left_shift", other, out=self)
+
+    def __irshift__(self, other):
+        return self._call_binary_np("right_shift", other, out=self)
+
+    def __iand__(self, other):
+        return self._call_binary("bitwise_and", other, out=self)
+
+    def __ior__(self, other):
+        return self._call_binary("bitwise_or", other, out=self)
+
+    def __ixor__(self, other):
+        return self._call_binary("bitwise_xor", other, out=self)
+
+    def copy(self, out=None):
+        return self._call_unary("copy", out=out)
+
+    def acos(self, out=None):
+        return self._call_unary("acos", out=out)
+
+    def acosh(self, out=None):
+        return self._call_unary("acosh", out=out)
+
+    def asin(self, out=None):
+        return self._call_unary("asin", out=out)
+
+    def asinh(self, out=None):
+        return self._call_unary("asinh", out=out)
+
+    def atan(self, out=None):
+        return self._call_unary("atan", out=out)
+
+    def atanh(self, out=None):
+        return self._call_unary("atanh", out=out)
+
+    def cbrt(self, out=None):
+        return self._call_unary("cbrt", out=out)
+
+    def cos(self, out=None):
+        return self._call_unary("cos", out=out)
+
+    def cosh(self, out=None):
+        return self._call_unary("cosh", out=out)
+
+    def erf(self, out=None):
+        return self._call_unary("erf", out=out)
+
+    def erfc(self, out=None):
+        return self._call_unary("erfc", out=out)
+
+    def exp(self, out=None):
+        return self._call_unary("exp", out=out)
+
+    def exp2(self, out=None):
+        return self._call_unary("exp2", out=out)
+
+    def expm1(self, out=None):
+        return self._call_unary("expm1", out=out)
+
+    def fabs(self, out=None):
+        return self._call_unary("fabs", out=out)
+
+    def lgamma(self, out=None):
+        return self._call_unary("lgamma", out=out)
+
+    def log(self, out=None):
+        return self._call_unary("log", out=out)
+
+    def log10(self, out=None):
+        return self._call_unary("log10", out=out)
+
+    def log1p(self, out=None):
+        return self._call_unary("log1p", out=out)
+
+    def log2(self, out=None):
+        return self._call_unary("log2", out=out)
+
+    def logb(self, out=None):
+        return self._call_unary("logb", out=out)
+
+    def nearbyint(self, out=None):
+        return self._call_unary("nearbyint", out=out)
+
+    def sin(self, out=None):
+        return self._call_unary("sin", out=out)
 
-    search(max_level, tree)
-    if acc[max_level] and all(x is None for x in acc[max_level]):
-        pass # min_level is not set in this special case, hence the check.
-    elif min_level != max_level:
-        raise ValueError("unbalanced tree: min depth: %d max depth: %d" %
-                         (min_level, max_level))
+    def sinh(self, out=None):
+        return self._call_unary("sinh", out=out)
 
-    data = acc[max_level]
-    shapes = list(reversed(acc[0:max_level]))
+    def sqrt(self, out=None):
+        return self._call_unary("sqrt", out=out)
 
-    return data, shapes
+    def tan(self, out=None):
+        return self._call_unary("tan", out=out)
 
+    def tanh(self, out=None):
+        return self._call_unary("tanh", out=out)
 
+    def tanh(self, out=None):
+        return self._call_unary("tgamma", out=out)
 
+    def equaln(self, other, out=None):
+        return self._call_binary("equaln", other, out=out, raiseit=True)