def __new__(cls,
value,
*,
type=None,
dtype=None,
levels=None,
typedef=None,
dtypedef=None):
if (type, dtype, levels, typedef, dtypedef).count(None) < 2:
raise TypeError(
"the 'type', 'dtype', 'levels' and 'typedef' arguments are "
"mutually exclusive")
if type is not None:
if isinstance(type, str):
type = ndt(type)
elif dtype is not None:
type = typeof(value, dtype=dtype)
elif levels is not None:
args = ', '.join("'%s'" % l if l is not None else 'NA'
for l in levels)
t = "%d * categorical(%s)" % (len(value), args)
type = ndt(t)
elif typedef is not None:
type = ndt(typedef)
if type.isabstract():
dtype = type.hidden_dtype
t = typeof(value, dtype=dtype)
type = instantiate(typedef, t)
elif dtypedef is not None:
dtype = ndt(dtypedef)
type = typeof(value, dtype=dtype)
else:
type = typeof(value)
return super().__new__(cls, type=type, value=value)
def test_multiple_return(self):
x, y = ex.randtuple()
self.assertEqual(x.type, ndt("int32"))
self.assertEqual(y.type, ndt("int32"))
x, y = ex.divmod10(xnd(233))
self.assertEqual(x.value, 23)
self.assertEqual(y.value, 3)
def test_const_incr(self):
@numba_xnd.gumath.register_kernel_direct("test_const_incr",
"int64 -> int64")
@numba_xnd.gumath.wrap_kernel_dispatcher(2)
@numba.njit((
numba_xnd.libxnd.xnd_view_t.WrapperNumbaType(ndt("int64")),
numba_xnd.libxnd.xnd_view_t.WrapperNumbaType(ndt("int64")),
))
def something(a, ret):
ret[()] = a.value + 10
self.assertEqual(something(xnd(10)), xnd(20))
def xndvectorize(v):
if isinstance(v, str):
v = [v]
if isinstance(v, list):
lst = [ndt(s).to_nbformat() for s in v]
sigs = [x[0] for x in lst]
coretypes = [x[1] for x in lst]
if (len(set(sigs))) != 1:
raise ValueError(
"empty list or different signatures in multimethod")
sig = sigs[0]
else:
raise TypeError("unsupported input type %s" % type(v))
def wrap(func):
guvec = GUVectorize(func, sig, nopython=True)
for t in coretypes:
guvec.add(t)
g = guvec.build_ufunc()
def h(*args, **kwargs):
out = g(*args, **kwargs)
view = xnd.from_buffer(out)
ret = xnd.empty(view.type)
np.copyto(np.array(ret, copy=False), out)
return ret
return h
return wrap
def from_buffer_and_type(cls, obj=None, type=None):
"""Return an xnd object that obtains memory from 'obj' via the
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 super().from_buffer_and_type(obj, type)
def __new__(cls, value, *, type=None, dtype=None, levels=None):
if (type, dtype, levels).count(None) < 2:
raise TypeError(
"the 'type', 'dtype' and 'levels' arguments are mutually "
"exclusive")
if type is not None:
if isinstance(type, str):
type = ndt(type)
elif dtype is not None:
type = typeof(value, dtype=dtype)
elif levels is not None:
args = ', '.join("'%s'" % l if l is not None else 'NA'
for l in levels)
t = "%d * categorical(%s)" % (len(value), args)
type = ndt(t)
else:
type = typeof(value)
return super().__new__(cls, type=type, value=value)
def unsafe_from_data(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.
"""
if isinstance(type, str):
type = ndt(type)
return cls._unsafe_from_data(obj, type)
def xnd_wrapper_value(x_wrapper):
n = x_wrapper.ndt_value
if n == ndtypes.ndt("int64"):
def get(x_wrapper):
x = unwrap_xnd_view(x_wrapper).view
return shared.ptr_load_type(numba.types.int64, x.ptr)
return get
if n == ndtypes.ndt("float64"):
def get(x_wrapper):
x = unwrap_xnd_view(x_wrapper).view
return shared.ptr_load_type(numba.types.float64, x.ptr)
return get
def __ua_convert__(value, dispatch_type, coerce):
if dispatch_type is ndarray:
return convert(value, coerce=coerce) if value is not None else None
if dispatch_type is ufunc and hasattr(fn, value.name):
return getattr(fn, value.name)
if dispatch_type is dtype:
return ndt(str(value)) if value is not None else None
return NotImplemented
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
def sig_to_stack(t):
"""
Takes in a ndtypes signature and returns a list of the types for the xnd objects
on the stack passed into the gumath kernel.
>>> sig_to_stack(ndtypes.ndt("... * int64, ... * N * int64 -> float64, int64"))
(ndtypes.ndt("int64"), ndtypes.ndt("N * int64"), ndtypes.ndt("float64"), ndtypes.ndt("int64"))
"""
args, rets = str(t).split(" -> ")
stacks = tuple()
for v in args.split(", ") + rets.split(", "):
stacks += (ndtypes.ndt(" * ".join(d for d in v.split(" * ")
if "..." not in d)), )
return stacks
def __ua_convert__(value, dispatch_type, coerce):
if dispatch_type is ufunc and hasattr(fn, value.name):
return getattr(fn, value.name)
if value is None:
return None
if dispatch_type is ndarray:
if not coerce and not isinstance(value, xnd.xnd):
return NotImplemented
return convert(value, coerce=coerce)
if dispatch_type is dtype:
return ndt(str(value))
return NotImplemented
def wrap_impl(self, typingctx, inner_t, ndt_type_t):
if inner_t != self.numba_type:
return
# supports passing in strings as ndt's
if isinstance(ndt_type_t, numba.types.Const):
n = ndtypes.ndt(ndt_type_t.value)
arg_type = numba.types.string
elif hasattr(ndt_type_t, "ndt_value"):
n = ndt_type_t.ndt_value
arg_type = ndt_type_t
else:
return
sig = self.WrapperNumbaType(n)(self.numba_type, arg_type)
def codegen(context, builder, sig, args):
return args[0]
return sig, codegen
def test_functions_coerce_with_dtype(backend, method, args, kwargs):
backend, types = backend
for dtype in dtypes:
try:
with ua.set_backend(backend, coerce=True):
kwargs["dtype"] = dtype
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend,
method) not in EXCEPTIONS:
raise
pytest.xfail(
reason="The backend has no implementation for this ufunc.")
assert isinstance(ret, types)
if XndBackend is not None and backend == XndBackend:
assert ret.dtype == ndt(dtype)
else:
assert ret.dtype == dtype
return None
def reduce(f, x, axes=0, dtype=None):
if dtype is None:
dtype = maxcast[x.dtype]
g = get_cuda_reduction_func(f)
if g is not None:
return reduce_cuda(g, x, axes, dtype)
return reduce_cpu(f, x, axes, dtype)
maxcast = {
ndt("int8"): ndt("int64"),
ndt("int16"): ndt("int64"),
ndt("int32"): ndt("int64"),
ndt("int64"): ndt("int64"),
ndt("uint8"): ndt("uint64"),
ndt("uint16"): ndt("uint64"),
ndt("uint32"): ndt("uint64"),
ndt("uint64"): ndt("uint64"),
ndt("bfloat16"): ndt("float64"),
ndt("float16"): ndt("float64"),
ndt("float32"): ndt("float64"),
ndt("float64"): ndt("float64"),
ndt("complex32"): ndt("complex128"),
ndt("complex64"): ndt("complex128"),
ndt("complex128"): ndt("complex128"),
ndt("?int8"): ndt("?int64"),
def ndtypes_index(t):
"""
Returns the resulting ndtype after indexing t by some int.
"""
first, *rest = str(t).split(" * ")
return ndtypes.ndt(" * ".join(rest))
def typeof(value, *, dtype=None):
return ndt(_typeof(value, dtype=dtype))
def typeof(value, *, dtype=None):
s = _typeof(value, dtype=dtype)
t = s.replace("Any", "float64")
return ndt(t)
def genindices(factor):
for i in range(4):
yield ()
for i in range(4):
yield (factor * i, )
for i in range(4):
for j in range(4):
yield (factor * i, factor * j)
for i in range(4):
for j in range(4):
for k in range(4):
yield (factor * i, factor * j, factor * k)
BROADCAST_TEST_CASES = [
dict(sig=ndt("uint8 -> float64"),
args=[ndt("uint8")],
out=None,
spec=ApplySpec(flags='C|Fortran|Strided|Xnd',
outer_dims=0,
nin=1,
nout=1,
nargs=2,
types=[ndt("uint8"), ndt("float64")])),
dict(sig=ndt("... * uint8 -> ... * float64"),
args=[ndt("2 * uint8")],
out=None,
spec=ApplySpec(flags='OptZ|OptC|OptS|C|Fortran|Strided|Xnd',
outer_dims=1,
nin=1,
nout=1,
import unittest
from ndtypes import ndt
import numba_xnd # NOQA
from numba import njit
n = ndt("10 * 4 * 4 * 5 * 10 * int64")
# class TestBoxingUnboxing(unittest.TestCase):
# def test_boxes_unboxes(self):
# self.assertEqual(njit(lambda x: x)(n), n)
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
if dtype is None:
if not data:
dtype = 'float64'
else:
dtype = _choose_dtype(data)
for x in data:
if x is not None:
t = _typeof(x)
if t != dtype:
dtype = str(ndt(t).unify(ndt(dtype)))
if opt:
dtype = '?' + dtype
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)
return t
elif dtype is not None:
raise TypeError("dtype argument is only supported for arrays")
elif isinstance(value, tuple):
return "(" + ", ".join([_typeof(x) for x in value]) + ")"
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")
elif value is None:
return '?Any'
elif isinstance(value, float):
return 'float64'
elif isinstance(value, complex):
return 'complex128'
elif isinstance(value, int):
return 'int64'
elif isinstance(value, str):
return 'string'
elif isinstance(value, bytes):
return 'bytes'
else:
raise ValueError("cannot infer type for %r" % value)
import ndtypes
import xnd_structinfo as xinfo
# print('\n'.join(sorted(dir(xinfo))))
n = ndtypes.ndt("2 * 3 * float64")
print(n)
cn = xinfo.capsulate_NdtObject(n)
ndt = xinfo.get_NdtObject_ndt(cn)
def show_ndt_members(ndt, tab=""):
tag = xinfo.value_int32(xinfo.get_ndt_t_tag(ndt))
access = xinfo.value_int32(xinfo.get_ndt_t_access(ndt))
flags = xinfo.value_int32(xinfo.get_ndt_t_flags(ndt))
ndim = xinfo.value_int32(xinfo.get_ndt_t_ndim(ndt))
datasize = xinfo.value_int64(xinfo.get_ndt_t_datasize(ndt))
if ndim > 0:
shape = xinfo.value_int64(xinfo.get_ndt_t_FixedDim_shape(ndt))
else:
shape = "NA"
print(
f"{tab}tag={tag}, access={access}, flags={flags}, ndim={ndim}, datasize={datasize}, shape={shape}"
)
if ndim > 0:
show_ndt_members(xinfo.get_ndt_t_FixedDim_type(ndt), tab=tab + " ")
def copy_contiguous(self, dtype=None):
if isinstance(dtype, str):
dtype = ndt(dtype)
return super().copy_contiguous(dtype=dtype)
for i in range(4):
yield ()
for i in range(4):
yield (factor * i, )
for i in range(4):
for j in range(4):
yield (factor * i, factor * j)
for i in range(4):
for j in range(4):
for k in range(4):
yield (factor * i, factor * j, factor * k)
BROADCAST_TEST_CASES = [
ApplySpec(tag='C',
sig=ndt("uint8 -> float64"),
in_types=[ndt("uint8")],
out_types=[ndt("float64")],
in_broadcast=[],
outer_dims=0),
ApplySpec(tag='C',
sig=ndt("... * uint8 -> ... * float64"),
in_types=[ndt("2 * uint8")],
out_types=[ndt("2 * float64")],
in_broadcast=[ndt("2 * uint8")],
outer_dims=1),
ApplySpec(tag='Fortran',
sig=ndt("... * uint8 -> ... * float64"),
in_types=[ndt("!2 * 3 * uint8")],
out_types=[ndt("!2 * 3 * float64")],
in_broadcast=[ndt("!2 * 3 * uint8")],
import unittest
from ndtypes import ndt
import numba_xnd
from numba import njit
n = ndt("10 * 4 * 4 * int64")
@njit
def get_ndim(x):
return numba_xnd.libndtypes.unwrap_ndt(x).ndim
@njit
def get_shape(x):
a = numba_xnd.libndtypes.create_ndt_ndarray()
numba_xnd.libndtypes.ndt_as_ndarray(
a, numba_xnd.libndtypes.unwrap_ndt(x),
numba_xnd.libndtypes.create_ndt_context())
return (a.shape[0], a.shape[1], a.shape[2])
@njit
def is_concrete(x):
return numba_xnd.libndtypes.ndt_is_concrete(
numba_xnd.libndtypes.unwrap_ndt(x))
class TestNdt(unittest.TestCase):
import itertools
from collections import Iterable
import numpy as np
import pandas as pd
import six
from pandas.api.types import is_array_like, is_bool_dtype, is_integer, is_integer_dtype
from pandas.core.arrays import ExtensionArray
from pandas.core.dtypes.dtypes import ExtensionDtype
import ndtypes
import xnd
from ndtypes import ndt
_python_type_map = {
str(ndt("int64").hidden_dtype): int,
str(ndt("?int64").hidden_dtype): int,
str(ndt("float64").hidden_dtype): float,
str(ndt("?float64").hidden_dtype): float,
str(ndt("string").hidden_dtype): six.text_type,
str(ndt("?string").hidden_dtype): six.text_type,
}
class XndframesDtype(ExtensionDtype):
def __init__(self, xnd_dtype):
self.xnd_dtype = xnd_dtype
def __str__(self):
return "xndframes[{}]".format(self.xnd_dtype)
def typeof(v, dtype=None):
if isinstance(dtype, str):
dtype = ndt(dtype)
return _typeof(v, dtype=dtype, shortcut=True)
def remove_outer_dimensions(t, n):
"""
Remove `n` outer dimensions from type `t`.
"""
return ndtypes.ndt(" * ".join(str(t).split(" * ")[n:]))
def test_void(self):
x = ex.randint()
self.assertEqual(x.type, ndt("int32"))
