`kwarray.arrayapi`¶

The ArrayAPI is a common API that works exactly the same on both torch.Tensors and numpy.ndarrays.

The ArrayAPI is a combination of efficiency and convinience. It is convinient because you can just use an operation directly, it will type check the data, and apply the appropriate method. But it is also efficient because it can be used with minimal type checking by accessing a type-specific backend.

For example, you can do:

impl = kwarray.ArrayAPI.coerce(data)

And then impl will give you direct access to the appropriate methods without any type checking overhead. e..g. impl.<op-you-want>(data)

But you can also do kwarray.ArrayAPI.<op-you-want>(data) on anything and it will do type checking and then do the operation you want.

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import torch
>>> import numpy as np
>>> data1 = torch.rand(10, 10)
>>> data2 = data1.numpy()
>>> # Method 1: grab the appropriate sub-impl
>>> impl1 = ArrayAPI.impl(data1)
>>> impl2 = ArrayAPI.impl(data2)
>>> result1 = impl1.sum(data1, axis=0)
>>> result2 = impl2.sum(data2, axis=0)
>>> res1_np = ArrayAPI.numpy(result1)
>>> res2_np = ArrayAPI.numpy(result2)
>>> print('res1_np = {!r}'.format(res1_np))
>>> print('res2_np = {!r}'.format(res2_np))
>>> assert np.allclose(res1_np, res2_np)
>>> # Method 2: choose the impl on the fly
>>> result1 = ArrayAPI.sum(data1, axis=0)
>>> result2 = ArrayAPI.sum(data2, axis=0)
>>> res1_np = ArrayAPI.numpy(result1)
>>> res2_np = ArrayAPI.numpy(result2)
>>> print('res1_np = {!r}'.format(res1_np))
>>> print('res2_np = {!r}'.format(res2_np))
>>> assert np.allclose(res1_np, res2_np)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import torch
>>> import numpy as np
>>> data1 = torch.rand(10, 10)
>>> data2 = data1.numpy()

Module Contents¶

Classes¶

`_ImplRegistry`
`TorchImpls`	Torch backend for the ArrayAPI API
`NumpyImpls`	Numpy backend for the ArrayAPI API
`ArrayAPI`	Compatability API between torch and numpy.

Functions¶

_get_funcname(func)

_torch_dtype_lut()

dtype_info(dtype)

Parameters: dtype (type) -- a numpy, torch, or python numeric data type

Attributes¶

`torch`
`_REGISTERY`
`_torchmethod`
`_numpymethod`
`_apimethod`
`TorchNumpyCompat`

kwarray.arrayapi.torch¶

kwarray.arrayapi._get_funcname(func)¶

class kwarray.arrayapi._ImplRegistry¶

Bases: object

_register(self, func, func_type, impl)¶

_implmethod(self, func=None, func_type='data_func', impl=None)¶

_apimethod(self, key=None, func_type='data_func')¶: Creates wrapper for a “data method” — i.e. a ArrayAPI function that has only one main argument, which is an array.

_ensure_datamethods_names_are_registered(self)¶: Checks to make sure all methods are implemented in both torch and numpy implementations as well as exposed in the ArrayAPI.

kwarray.arrayapi._REGISTERY¶

kwarray.arrayapi._torchmethod¶

kwarray.arrayapi._numpymethod¶

kwarray.arrayapi._apimethod¶

class kwarray.arrayapi.TorchImpls¶

Bases: object

Torch backend for the ArrayAPI API

is_tensor = True¶

is_numpy = False¶

ensure¶

cat(datas, axis=- 1)¶

hstack(datas)¶

Concatenates along axis=0 if inputs are are 1-D otherwise axis=1

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> datas1 = [torch.arange(10), torch.arange(10)]
>>> datas2 = [d.numpy() for d in datas1]
>>> ans1 = TorchImpls.hstack(datas1)
>>> ans2 = NumpyImpls.hstack(datas2)
>>> assert np.all(ans1.numpy() == ans2)

vstack(datas)¶

Ensures that inputs datas are at least 2D (prepending a dimension of 1) and then concats along axis=0.

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> datas1 = [torch.arange(10), torch.arange(10)]
>>> datas2 = [d.numpy() for d in datas1]
>>> ans1 = TorchImpls.vstack(datas1)
>>> ans2 = NumpyImpls.vstack(datas2)
>>> assert np.all(ans1.numpy() == ans2)

atleast_nd(arr, n, front=False)¶

view(data, *shape)¶

take(data, indices, axis=None)¶

compress(data, flags, axis=None)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> data = torch.rand(10, 4, 2)
>>> impl = kwarray.ArrayAPI.coerce(data)

>>> axis = 0
>>> flags = (torch.arange(data.shape[axis]) % 2) == 0
>>> out = impl.compress(data, flags, axis=axis)
>>> assert tuple(out.shape) == (5, 4, 2)

>>> axis = 1
>>> flags = (torch.arange(data.shape[axis]) % 2) == 0
>>> out = impl.compress(data, flags, axis=axis)
>>> assert tuple(out.shape) == (10, 2, 2)

>>> axis = 2
>>> flags = (torch.arange(data.shape[axis]) % 2) == 0
>>> out = impl.compress(data, flags, axis=axis)
>>> assert tuple(out.shape) == (10, 4, 1)

>>> axis = None
>>> data = torch.rand(10)
>>> flags = (torch.arange(data.shape[0]) % 2) == 0
>>> out = impl.compress(data, flags, axis=axis)
>>> assert tuple(out.shape) == (5,)

tile(data, reps)¶

Implement np.tile in torch

Example

>>> # xdoctest: +SKIP
>>> # xdoctest: +REQUIRES(module:torch)
>>> data = torch.arange(10)[:, None]
>>> ans1 = ArrayAPI.tile(data, [1, 2])
>>> ans2 = ArrayAPI.tile(data.numpy(), [1, 2])
>>> assert np.all(ans1.numpy() == ans2)

Doctest:

>>> # xdoctest: +SKIP
>>> # xdoctest: +REQUIRES(module:torch)
>>> shapes = [(3,), (3, 4,), (3, 5, 7), (1,), (3, 1, 3)]
>>> for shape in shapes:
>>>     data = torch.rand(*shape)
>>>     for axis in range(len(shape)):
>>>         for reps in it.product(*[range(0, 4)] * len(shape)):
>>>             ans1 = ArrayAPI.tile(data, reps)
>>>             ans2 = ArrayAPI.tile(data.numpy(), reps)
>>>             #print('ans1.shape = {!r}'.format(tuple(ans1.shape)))
>>>             #print('ans2.shape = {!r}'.format(tuple(ans2.shape)))
>>>             assert np.all(ans1.numpy() == ans2)

abstract repeat(data, repeats, axis=None)¶

I’m not actually sure how to implement this efficiently

Example

>>> # xdoctest: +SKIP
>>> data = torch.arange(10)[:, None]
>>> ans1 = ArrayAPI.repeat(data, 2, axis=1)
>>> ans2 = ArrayAPI.repeat(data.numpy(), 2, axis=1)
>>> assert np.all(ans1.numpy() == ans2)

Doctest:

>>> # xdoctest: +SKIP
>>> shapes = [(3,), (3, 4,), (3, 5, 7)]
>>> for shape in shapes:
>>>     data = torch.rand(*shape)
>>>     for axis in range(len(shape)):
>>>         for repeats in range(0, 4):
>>>             ans1 = ArrayAPI.repeat(data, repeats, axis=axis)
>>>             ans2 = ArrayAPI.repeat(data.numpy(), repeats, axis=axis)
>>>             assert np.all(ans1.numpy() == ans2)

ArrayAPI.repeat(data, 2, axis=0) ArrayAPI.repeat(data.numpy(), 2, axis=0)

x = np.array([[1,2],[3,4]]) np.repeat(x, [1, 2], axis=0)

ArrayAPI.repeat(data.numpy(), [1, 2])

T(data)¶

transpose(data, axes)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> data1 = torch.rand(2, 3, 5)
>>> data2 = data1.numpy()
>>> res1 = ArrayAPI.transpose(data1, (2, 0, 1))
>>> res2 = ArrayAPI.transpose(data2, (2, 0, 1))
>>> assert np.all(res1.numpy() == res2)

numel(data)¶

full_like(data, fill_value, dtype=None)¶

empty_like(data, dtype=None)¶

zeros_like(data, dtype=None)¶

ones_like(data, dtype=None)¶

full(shape, fill_value, dtype=float)¶

empty(shape, dtype=float)¶

zeros(shape, dtype=float)¶

ones(shape, dtype=float)¶

argmax(data, axis=None)¶

argsort(data, axis=- 1, descending=False)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> rng = np.random.RandomState(0)
>>> data2 = rng.rand(5, 5)
>>> data1 = torch.from_numpy(data2)
>>> res1 = ArrayAPI.argsort(data1)
>>> res2 = ArrayAPI.argsort(data2)
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.argsort(data1, axis=1)
>>> res2 = ArrayAPI.argsort(data2, axis=1)
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.argsort(data1, axis=1, descending=True)
>>> res2 = ArrayAPI.argsort(data2, axis=1, descending=True)
>>> assert np.all(res1.numpy() == res2)
>>> data2 = rng.rand(5)
>>> data1 = torch.from_numpy(data2)
>>> res1 = ArrayAPI.argsort(data1, axis=0, descending=True)
>>> res2 = ArrayAPI.argsort(data2, axis=0, descending=True)
>>> assert np.all(res1.numpy() == res2)

max(data, axis=None)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> data1 = torch.rand(5, 5, 5, 5, 5, 5)
>>> data2 = data1.numpy()
>>> res1 = ArrayAPI.max(data1)
>>> res2 = ArrayAPI.max(data2)
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.max(data1, axis=(4, 0, 1))
>>> res2 = ArrayAPI.max(data2, axis=(4, 0, 1))
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.max(data1, axis=(5, -2))
>>> res2 = ArrayAPI.max(data2, axis=(5, -2))
>>> assert np.all(res1.numpy() == res2)

max_argmax(data, axis=None)¶: Note: this isn’t always gaurenteed to be compatibile with numpy if there are equal elements in data. See: >>> np.ones(10).argmax() # xdoctest: +IGNORE_WANT 0 >>> torch.ones(10).argmax() # xdoctest: +IGNORE_WANT tensor(9)

maximum(data1, data2, out=None)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> data1 = torch.rand(5, 5)
>>> data2 = torch.rand(5, 5)
>>> result1 = TorchImpls.maximum(data1, data2)
>>> result2 = NumpyImpls.maximum(data1.numpy(), data2.numpy())
>>> assert np.allclose(result1.numpy(), result2)

minimum(data1, data2, out=None)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> data1 = torch.rand(5, 5)
>>> data2 = torch.rand(5, 5)
>>> result1 = TorchImpls.minimum(data1, data2)
>>> result2 = NumpyImpls.minimum(data1.numpy(), data2.numpy())
>>> assert np.allclose(result1.numpy(), result2)

sum(data, axis=None)¶

nan_to_num(x, copy=True)¶

copy(data)¶

nonzero(data)¶

astype(data, dtype, copy=True)¶

tensor(data, device=ub.NoParam)¶

numpy(data)¶

tolist(data)¶

contiguous(data)¶

pad(data, pad_width, mode='constant')¶

asarray(data, dtype=None)¶

Cast data into a tensor representation

Example

>>> data = np.empty((2, 0, 196, 196), dtype=np.float32)

dtype_kind(data)¶: returns the numpy code for the data type kind

floor(data, out=None)¶

ceil(data, out=None)¶

ifloor(data, out=None)¶

iceil(data, out=None)¶

round(data, decimals=0, out=None)¶

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> rng = kwarray.ensure_rng(0)
>>> np_data = rng.rand(10) * 100
>>> pt_data = torch.from_numpy(np_data)
>>> a = kwarray.ArrayAPI.round(np_data)
>>> b = kwarray.ArrayAPI.round(pt_data)
>>> assert np.all(a == b.numpy())
>>> a = kwarray.ArrayAPI.round(np_data, 2)
>>> b = kwarray.ArrayAPI.round(pt_data, 2)
>>> assert np.all(a == b.numpy())

iround(data, out=None, dtype=int)¶

clip(data, a_min=None, a_max=None, out=None)¶

softmax(data, axis=None)¶

class kwarray.arrayapi.NumpyImpls¶

Bases: object

Numpy backend for the ArrayAPI API

is_tensor = False¶

is_numpy = True¶

hstack¶

vstack¶

matmul¶

nan_to_num¶

log¶

log2¶

any¶

all¶

copy¶

nonzero¶

ensure¶

clip¶

cat(datas, axis=- 1)¶

atleast_nd(arr, n, front=False)¶

view(data, *shape)¶

take(data, indices, axis=None)¶

compress(data, flags, axis=None)¶

repeat(data, repeats, axis=None)¶

tile(data, reps)¶

T(data)¶

transpose(data, axes)¶

numel(data)¶

empty_like(data, dtype=None)¶

full_like(data, fill_value, dtype=None)¶

zeros_like(data, dtype=None)¶

ones_like(data, dtype=None)¶

full(shape, fill_value, dtype=float)¶

empty(shape, dtype=float)¶

zeros(shape, dtype=float)¶

ones(shape, dtype=float)¶

argmax(data, axis=None)¶

argsort(data, axis=- 1, descending=False)¶

max(data, axis=None)¶

max_argmax(data, axis=None)¶

sum(data, axis=None)¶

maximum(data1, data2, out=None)¶

minimum(data1, data2, out=None)¶

astype(data, dtype, copy=True)¶

tensor(data, device=ub.NoParam)¶

numpy(data)¶

tolist(data)¶

contiguous(data)¶

pad(data, pad_width, mode='constant')¶

asarray(data, dtype=None)¶: Cast data into a numpy representation

dtype_kind(data)¶

floor(data, out=None)¶

ceil(data, out=None)¶

ifloor(data, out=None)¶

iceil(data, out=None)¶

round(data, decimals=0, out=None)¶

iround(data, out=None, dtype=int)¶

softmax(data, axis=None)¶

class kwarray.arrayapi.ArrayAPI¶

Bases: object

Compatability API between torch and numpy.

The API defines classmethods that work on both Tensors and ndarrays. As such the user can simply use kwarray.ArrayAPI.<funcname> and it will return the expected result for both Tensor and ndarray types.

However, this is inefficient because it requires us to check the type of the input for every API call. Therefore it is recommended that you use the ArrayAPI.coerce() function, which takes as input the data you want to operate on. It performs the type check once, and then returns another object that defines with an identical API, but specific to the given data type. This means that we can ignore type checks on future calls of the specific implementation. See examples for more details.

Example

>>> # Use the easy-to-use, but inefficient array api
>>> # xdoctest: +REQUIRES(module:torch)
>>> take = ArrayAPI.take
>>> np_data = np.arange(0, 143).reshape(11, 13)
>>> pt_data = torch.LongTensor(np_data)
>>> indices = [1, 3, 5, 7, 11, 13, 17, 21]
>>> idxs0 = [1, 3, 5, 7]
>>> idxs1 = [1, 3, 5, 7, 11]
>>> assert np.allclose(take(np_data, indices), take(pt_data, indices))
>>> assert np.allclose(take(np_data, idxs0, 0), take(pt_data, idxs0, 0))
>>> assert np.allclose(take(np_data, idxs1, 1), take(pt_data, idxs1, 1))

Example

>>> # Use the easy-to-use, but inefficient array api
>>> # xdoctest: +REQUIRES(module:torch)
>>> compress = ArrayAPI.compress
>>> np_data = np.arange(0, 143).reshape(11, 13)
>>> pt_data = torch.LongTensor(np_data)
>>> flags = (np_data % 2 == 0).ravel()
>>> f0 = (np_data % 2 == 0)[:, 0]
>>> f1 = (np_data % 2 == 0)[0, :]
>>> assert np.allclose(compress(np_data, flags), compress(pt_data, flags))
>>> assert np.allclose(compress(np_data, f0, 0), compress(pt_data, f0, 0))
>>> assert np.allclose(compress(np_data, f1, 1), compress(pt_data, f1, 1))

Example

>>> # Use ArrayAPI to coerce an identical API that doesnt do type checks
>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> np_data = np.arange(0, 15).reshape(3, 5)
>>> pt_data = torch.LongTensor(np_data)
>>> # The new ``impl`` object has the same API as ArrayAPI, but works
>>> # specifically on torch Tensors.
>>> impl = kwarray.ArrayAPI.coerce(pt_data)
>>> flat_data = impl.view(pt_data, -1)
>>> print('flat_data = {!r}'.format(flat_data))
flat_data = tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14])
>>> # The new ``impl`` object has the same API as ArrayAPI, but works
>>> # specifically on numpy ndarrays.
>>> impl = kwarray.ArrayAPI.coerce(np_data)
>>> flat_data = impl.view(np_data, -1)
>>> print('flat_data = {!r}'.format(flat_data))
flat_data = array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14])

_torch¶

_numpy¶

take¶

compress¶

repeat¶

tile¶

view¶

numel¶

atleast_nd¶

full_like¶

ones_like¶

zeros_like¶

empty_like¶

sum¶

argmax¶

argsort¶

max¶

maximum¶

minimum¶

matmul¶

astype¶

nonzero¶

nan_to_num¶

tensor¶

numpy¶

tolist¶

asarray¶

asarray¶

T¶

transpose¶

contiguous¶

pad¶

dtype_kind¶

max_argmax¶

any¶

all¶

log2¶

log¶

copy¶

iceil¶

ifloor¶

floor¶

ceil¶

round¶

iround¶

clip¶

softmax¶

static impl(data)¶

Returns a namespace suitable for operating on the input data type

Parameters: data (ndarray | Tensor) – data to be operated on

static coerce(data)¶: Coerces some form of inputs into an array api (either numpy or torch).

cat(datas, *args, **kwargs)¶

hstack(datas, *args, **kwargs)¶

vstack(datas, *args, **kwargs)¶

kwarray.arrayapi.TorchNumpyCompat¶

kwarray.arrayapi._torch_dtype_lut()¶

kwarray.arrayapi.dtype_info(dtype)¶

Parameters: dtype (type) – a numpy, torch, or python numeric data type
Returns: an iinfo of finfo structure depending on the input type.
Return type: struct

References

https://higra.readthedocs.io/en/stable/_modules/higra/hg_utils.html#dtype_info

Example

>>> from kwarray.arrayapi import *  # NOQA
>>> results = []
>>> results += [dtype_info(float)]
>>> results += [dtype_info(int)]
>>> results += [dtype_info(complex)]
>>> results += [dtype_info(np.float32)]
>>> results += [dtype_info(np.int32)]
>>> results += [dtype_info(np.uint32)]
>>> if hasattr(np, 'complex256'):
>>>     results += [dtype_info(np.complex256)]
>>> if torch is not None:
>>>     results += [dtype_info(torch.float32)]
>>>     results += [dtype_info(torch.int64)]
>>>     results += [dtype_info(torch.complex64)]
>>> for info in results:
>>>     print('info = {!r}'.format(info))
>>> for info in results:
>>>     print('info.bits = {!r}'.format(info.bits))

kwarray.arrayapi¶

Module Contents¶

Classes¶

Functions¶

Attributes¶

`kwarray.arrayapi`¶