kwarray.arrayapi module

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.

Idea:

Perhaps we could separate this into its own python package (maybe called “onearray”), where the module itself behaves like the ArrayAPI. Design goals are to provide easy to use (as drop-in as possible) replacements for torch or numpy function calls. It has to have near-zero overhead, or at least a way to make that happen.

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> import torch
>>> import numpy as np
>>> data1 = torch.rand(10, 10)
>>> data2 = data1.numpy()
>>> # Method 1: grab the appropriate sub-impl
>>> impl1 = kwarray.ArrayAPI.impl(data1)
>>> impl2 = kwarray.ArrayAPI.impl(data2)
>>> result1 = impl1.sum(data1, axis=0)
>>> result2 = impl2.sum(data2, axis=0)
>>> res1_np = kwarray.ArrayAPI.numpy(result1)
>>> res2_np = kwarray.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 = kwarray.ArrayAPI.sum(data1, axis=0)
>>> result2 = kwarray.ArrayAPI.sum(data2, axis=0)
>>> res1_np = kwarray.ArrayAPI.numpy(result1)
>>> res2_np = kwarray.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()

class kwarray.arrayapi._ImplRegistry

Bases: object

_register(func, func_type, impl)

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

_apimethod(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()

Checks to make sure all methods are implemented in both torch and numpy implementations as well as exposed in the ArrayAPI.

kwarray.arrayapi._torchmethod(func=None, func_type='data_func', *, impl='torch')

kwarray.arrayapi._numpymethod(func=None, func_type='data_func', *, impl='numpy')

kwarray.arrayapi._apimethod(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.

class kwarray.arrayapi.TorchImpls

Bases: object

Torch backend for the ArrayAPI API

is_tensor = True

is_numpy = False

static result_type(*arrays_and_dtypes)

Return type from promotion rules

dtype, promote_types, min_scalar_type, can_cast

static cat(datas, axis=-1)

static hstack(datas)

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

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import 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)

static 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)
>>> import 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)

static atleast_nd(arr, n, front=False)

static view(data, *shape)

static take(data, indices, axis=None)

static compress(data, flags, axis=None)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> import torch
>>> 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,)

static 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)

static 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])

static T(data)

static transpose(data, axes)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import 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)

static numel(data)

static full_like(data, fill_value, dtype=None)

static empty_like(data, dtype=None)

static zeros_like(data, dtype=None)

static ones_like(data, dtype=None)

static full(shape, fill_value, dtype=<class 'float'>)

static empty(shape, dtype=<class 'float'>)

static zeros(shape, dtype=<class 'float'>)

static ones(shape, dtype=<class 'float'>)

static argmax(data, axis=None)

static argmin(data, axis=None)

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

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> import torch
>>> 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)

static max(data, axis=None)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> from kwarray.arrayapi import _TORCH_HAS_MAX_BUG
>>> import torch
>>> import pytest
>>> if _TORCH_HAS_MAX_BUG:
>>>     pytest.skip('torch max has a bug, which was fixed in 1.7')
>>> 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)

static min(data, axis=None)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> from kwarray.arrayapi import _TORCH_HAS_MAX_BUG
>>> import pytest
>>> import torch
>>> if _TORCH_HAS_MAX_BUG:
>>>     pytest.skip('torch min has a bug, which was fixed in 1.7')
>>> data1 = torch.rand(5, 5, 5, 5, 5, 5)
>>> data2 = data1.numpy()
>>> res1 = ArrayAPI.min(data1)
>>> res2 = ArrayAPI.min(data2)
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.min(data1, axis=(4, 0, 1))
>>> res2 = ArrayAPI.min(data2, axis=(4, 0, 1))
>>> assert np.all(res1.numpy() == res2)
>>> res1 = ArrayAPI.min(data1, axis=(5, -2))
>>> res2 = ArrayAPI.min(data2, axis=(5, -2))
>>> assert np.all(res1.numpy() == res2)

static max_argmax(data, axis=None)

Parameters:

• data (Tensor) – data to perform operation on

• axis (None | int) – axis to perform operation on

Returns:

The max value(s) and argmax position(s).

Return type:

Tuple[Numeric, Numeric]

Note

In modern versions of torch and numpy if there are multiple maximum values the index of the instance is returned. This is not true in older versions of torch. I’m unsure when this gaurentee was added to numpy.

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> from kwarray.arrayapi import ArrayAPI
>>> import torch
>>> data = 1 / (1 + (torch.arange(12) - 6).view(3, 4) ** 2)
>>> ArrayAPI.max_argmax(data)
(tensor(1...), tensor(6))
>>> # When the values are all the same, there doesn't seem
>>> # to be a reliable spec on which one is returned first.
>>> np.ones(10).argmax()   # xdoctest: +IGNORE_WANT
0
>>> # Newer versions of torch (e.g. 1.12.0)
>>> torch.ones(10).argmax()   # xdoctest: +IGNORE_WANT
tensor(0)
>>> # Older versions of torch  (e.g 1.6.0)
>>> torch.ones(10).argmax()   # xdoctest: +IGNORE_WANT
tensor(9)

static min_argmin(data, axis=None)

Parameters:

• data (Tensor) – data to perform operation on

• axis (None | int) – axis to perform operation on

Returns:

The min value(s) and argmin position(s).

Return type:

Tuple[Numeric, Numeric]

Note

In modern versions of torch and numpy if there are multiple minimum values the index of the instance is returned. This is not true in older versions of torch. I’m unsure when this gaurentee was added to numpy.

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> from kwarray.arrayapi import ArrayAPI
>>> import torch
>>> data = (torch.arange(12) - 6).view(3, 4) ** 2
>>> ArrayAPI.min_argmin(data)
(tensor(0), tensor(6))
>>> # Issue demo:
>>> # When the values are all the same, there doesn't seem
>>> # to be a reliable spec on which one is returned first.
>>> np.ones(10).argmin()   # xdoctest: +IGNORE_WANT
0
>>> # Newer versions of torch (e.g. 1.12.0)
>>> torch.ones(10).argmin()   # xdoctest: +IGNORE_WANT
tensor(0)
>>> # Older versions of torch  (e.g 1.6.0)
>>> torch.ones(10).argmin()   # xdoctest: +IGNORE_WANT
tensor(9)

static maximum(data1, data2, out=None)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import sys, ubelt
>>> from kwarray.arrayapi import *  # NOQA
>>> import 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)
>>> result1 = TorchImpls.maximum(data1, 0)
>>> result2 = NumpyImpls.maximum(data1.numpy(), 0)
>>> assert np.allclose(result1.numpy(), result2)

static minimum(data1, data2, out=None)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import 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)

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import sys, ubelt
>>> from kwarray.arrayapi import *  # NOQA
>>> import 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)
>>> result1 = TorchImpls.minimum(data1, 0)
>>> result2 = NumpyImpls.minimum(data1.numpy(), 0)
>>> assert np.allclose(result1.numpy(), result2)

static array_equal(data1, data2, equal_nan=False) → bool

Returns True if all elements in the array are equal. This mirrors the behavior of numpy.array_equal().

Parameters:

• data1 (Tensor) – first array

• data2 (Tensor) – second array

• equal_nan (bool) – Whether to compare NaN’s as equal.

Returns:

bool

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> import torch
>>> data1 = torch.rand(5, 5)
>>> data2 = data1 + 1
>>> result1 = TorchImpls.array_equal(data1, data2)
>>> result2 = NumpyImpls.array_equal(data1.numpy(), data2.numpy())
>>> result3 = TorchImpls.array_equal(data1, data1)
>>> result4 = NumpyImpls.array_equal(data1.numpy(), data1.numpy())
>>> assert result1 is False
>>> assert result2 is False
>>> assert result3 is True
>>> assert result4 is True

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> from kwarray.arrayapi import *  # NOQA
>>> import torch
>>> data1 = torch.rand(5, 5)
>>> data1[0] = np.nan
>>> data2 = data1
>>> result1 = TorchImpls.array_equal(data1, data2)
>>> result2 = NumpyImpls.array_equal(data1.numpy(), data2.numpy())
>>> result3 = TorchImpls.array_equal(data1, data2, equal_nan=True)
>>> result4 = NumpyImpls.array_equal(data1.numpy(), data2.numpy(), equal_nan=True)
>>> assert result1 is False
>>> assert result2 is False
>>> assert result3 is True
>>> assert result4 is True

static matmul(data1, data2, out=None)

static sum(data, axis=None)

static nan_to_num(x, copy=True)

static copy(data)

static log(data)

static log2(data)

static any(data)

static all(data)

static nonzero(data)

static astype(data, dtype, copy=True)

static tensor(data, device=NoParam)

static numpy(data)

static tolist(data)

static contiguous(data)

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

static asarray(data, dtype=None)

Cast data into a tensor representation

Example

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

static ensure(data, dtype=None)

Cast data into a tensor representation

Example

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

static dtype_kind(data)

returns the numpy code for the data type kind

static floor(data, out=None)

static ceil(data, out=None)

static ifloor(data, out=None)

static iceil(data, out=None)

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

Example

>>> # xdoctest: +REQUIRES(module:torch)
>>> import kwarray
>>> import torch
>>> 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())

static iround(data, out=None, dtype=<class 'int'>)

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

static softmax(data, axis=None)

class kwarray.arrayapi.NumpyImpls

Bases: object

Numpy backend for the ArrayAPI API

is_tensor = False

is_numpy = True

static hstack(tup, *, dtype=None, casting='same_kind')

Stack arrays in sequence horizontally (column wise).

This is equivalent to concatenation along the second axis, except for 1-D arrays where it concatenates along the first axis. Rebuilds arrays divided by hsplit.

This function makes most sense for arrays with up to 3 dimensions. For instance, for pixel-data with a height (first axis), width (second axis), and r/g/b channels (third axis). The functions concatenate, stack and block provide more general stacking and concatenation operations.

Parameters:

• tup (sequence of ndarrays) – The arrays must have the same shape along all but the second axis, except 1-D arrays which can be any length.

• dtype (str or dtype) – If provided, the destination array will have this dtype. Cannot be provided together with out.

• .. versionadded:: 1.24

• casting ({‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}, optional) – Controls what kind of data casting may occur. Defaults to ‘same_kind’.

• .. versionadded:: 1.24

Returns:

stacked – The array formed by stacking the given arrays.

Return type:

ndarray

See also

concatenate

Join a sequence of arrays along an existing axis.

stack

Join a sequence of arrays along a new axis.

block

Assemble an nd-array from nested lists of blocks.

vstack

Stack arrays in sequence vertically (row wise).

dstack

Stack arrays in sequence depth wise (along third axis).

column_stack

Stack 1-D arrays as columns into a 2-D array.

hsplit

Split an array into multiple sub-arrays horizontally (column-wise).

Examples

>>> a = np.array((1,2,3))
>>> b = np.array((4,5,6))
>>> np.hstack((a,b))
array([1, 2, 3, 4, 5, 6])
>>> a = np.array([[1],[2],[3]])
>>> b = np.array([[4],[5],[6]])
>>> np.hstack((a,b))
array([[1, 4],
       [2, 5],
       [3, 6]])

static vstack(tup, *, dtype=None, casting='same_kind')

Stack arrays in sequence vertically (row wise).

This is equivalent to concatenation along the first axis after 1-D arrays of shape (N,) have been reshaped to (1,N). Rebuilds arrays divided by vsplit.

This function makes most sense for arrays with up to 3 dimensions. For instance, for pixel-data with a height (first axis), width (second axis), and r/g/b channels (third axis). The functions concatenate, stack and block provide more general stacking and concatenation operations.

np.row_stack is an alias for vstack. They are the same function.

Parameters:

• tup (sequence of ndarrays) – The arrays must have the same shape along all but the first axis. 1-D arrays must have the same length.

• dtype (str or dtype) – If provided, the destination array will have this dtype. Cannot be provided together with out.

• .. versionadded:: 1.24

• casting ({‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}, optional) – Controls what kind of data casting may occur. Defaults to ‘same_kind’.

• .. versionadded:: 1.24

Returns:

stacked – The array formed by stacking the given arrays, will be at least 2-D.

Return type:

ndarray

See also

concatenate

Join a sequence of arrays along an existing axis.

stack

Join a sequence of arrays along a new axis.

block

Assemble an nd-array from nested lists of blocks.

hstack

Stack arrays in sequence horizontally (column wise).

dstack

Stack arrays in sequence depth wise (along third axis).

column_stack

Stack 1-D arrays as columns into a 2-D array.

vsplit

Split an array into multiple sub-arrays vertically (row-wise).

Examples

>>> a = np.array([1, 2, 3])
>>> b = np.array([4, 5, 6])
>>> np.vstack((a,b))
array([[1, 2, 3],
       [4, 5, 6]])

>>> a = np.array([[1], [2], [3]])
>>> b = np.array([[4], [5], [6]])
>>> np.vstack((a,b))
array([[1],
       [2],
       [3],
       [4],
       [5],
       [6]])

static result_type(*arrays_and_dtypes)

Return type from promotion rules

SeeAlso:

numpy.find_common_type() numpy.promote_types() numpy.result_type()

static cat(datas, axis=-1)

static atleast_nd(arr, n, front=False)

static view(data, *shape)

static take(data, indices, axis=None)

static compress(data, flags, axis=None)

static repeat(data, repeats, axis=None)

static tile(data, reps)

static T(data)

static transpose(data, axes)

static numel(data)

static empty_like(data, dtype=None)

static full_like(data, fill_value, dtype=None)

static zeros_like(data, dtype=None)

static ones_like(data, dtype=None)

static full(shape, fill_value, dtype=<class 'float'>)

static empty(shape, dtype=<class 'float'>)

static zeros(shape, dtype=<class 'float'>)

static ones(shape, dtype=<class 'float'>)

static argmax(data, axis=None)

static argmin(data, axis=None)

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

static max(data, axis=None)

static min(data, axis=None)

static max_argmax(data, axis=None)

Parameters:

• data (ArrayLike) – data to perform operation on

• axis (None | int) – axis to perform operation on

Returns:

The max value(s) and argmax position(s).

Return type:

Tuple[Numeric, Numeric]

static min_argmin(data, axis=None)

Parameters:

• data (ArrayLike) – data to perform operation on

• axis (None | int) – axis to perform operation on

Returns:

The min value(s) and argmin position(s).

Return type:

Tuple[Numeric, Numeric]

static sum(data, axis=None)

static maximum(data1, data2, out=None)

static minimum(data1, data2, out=None)

matmul = <ufunc 'matmul'>

static nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None)

Replace NaN with zero and infinity with large finite numbers (default behaviour) or with the numbers defined by the user using the nan, posinf and/or neginf keywords.

If x is inexact, NaN is replaced by zero or by the user defined value in nan keyword, infinity is replaced by the largest finite floating point values representable by x.dtype or by the user defined value in posinf keyword and -infinity is replaced by the most negative finite floating point values representable by x.dtype or by the user defined value in neginf keyword.

For complex dtypes, the above is applied to each of the real and imaginary components of x separately.

If x is not inexact, then no replacements are made.

Parameters:

• x (scalar or array_like) – Input data.

• copy (bool, optional) – Whether to create a copy of x (True) or to replace values in-place (False). The in-place operation only occurs if casting to an array does not require a copy. Default is True.

  New in version 1.13.

• nan (int, float, optional) – Value to be used to fill NaN values. If no value is passed then NaN values will be replaced with 0.0.

  New in version 1.17.

• posinf (int, float, optional) – Value to be used to fill positive infinity values. If no value is passed then positive infinity values will be replaced with a very large number.

  New in version 1.17.

• neginf (int, float, optional) – Value to be used to fill negative infinity values. If no value is passed then negative infinity values will be replaced with a very small (or negative) number.

  New in version 1.17.

Returns:

out – x, with the non-finite values replaced. If copy is False, this may be x itself.

Return type:

ndarray

See also

isinf

Shows which elements are positive or negative infinity.

isneginf

Shows which elements are negative infinity.

isposinf

Shows which elements are positive infinity.

isnan

Shows which elements are Not a Number (NaN).

isfinite

Shows which elements are finite (not NaN, not infinity)

Notes

NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic (IEEE 754). This means that Not a Number is not equivalent to infinity.

Examples

>>> np.nan_to_num(np.inf)
1.7976931348623157e+308
>>> np.nan_to_num(-np.inf)
-1.7976931348623157e+308
>>> np.nan_to_num(np.nan)
0.0
>>> x = np.array([np.inf, -np.inf, np.nan, -128, 128])
>>> np.nan_to_num(x)
array([ 1.79769313e+308, -1.79769313e+308,  0.00000000e+000, # may vary
       -1.28000000e+002,  1.28000000e+002])
>>> np.nan_to_num(x, nan=-9999, posinf=33333333, neginf=33333333)
array([ 3.3333333e+07,  3.3333333e+07, -9.9990000e+03,
       -1.2800000e+02,  1.2800000e+02])
>>> y = np.array([complex(np.inf, np.nan), np.nan, complex(np.nan, np.inf)])
array([  1.79769313e+308,  -1.79769313e+308,   0.00000000e+000, # may vary
     -1.28000000e+002,   1.28000000e+002])
>>> np.nan_to_num(y)
array([  1.79769313e+308 +0.00000000e+000j, # may vary
         0.00000000e+000 +0.00000000e+000j,
         0.00000000e+000 +1.79769313e+308j])
>>> np.nan_to_num(y, nan=111111, posinf=222222)
array([222222.+111111.j, 111111.     +0.j, 111111.+222222.j])

static array_equal(a1, a2, equal_nan=False)

True if two arrays have the same shape and elements, False otherwise.

Parameters:

• a1, a2 (array_like) – Input arrays.

• equal_nan (bool) – Whether to compare NaN’s as equal. If the dtype of a1 and a2 is complex, values will be considered equal if either the real or the imaginary component of a given value is nan.

  New in version 1.19.0.

Returns:

b – Returns True if the arrays are equal.

Return type:

bool

See also

allclose

Returns True if two arrays are element-wise equal within a tolerance.

array_equiv

Returns True if input arrays are shape consistent and all elements equal.

Examples

>>> np.array_equal([1, 2], [1, 2])
True
>>> np.array_equal(np.array([1, 2]), np.array([1, 2]))
True
>>> np.array_equal([1, 2], [1, 2, 3])
False
>>> np.array_equal([1, 2], [1, 4])
False
>>> a = np.array([1, np.nan])
>>> np.array_equal(a, a)
False
>>> np.array_equal(a, a, equal_nan=True)
True

When equal_nan is True, complex values with nan components are considered equal if either the real or the imaginary components are nan.

>>> a = np.array([1 + 1j])
>>> b = a.copy()
>>> a.real = np.nan
>>> b.imag = np.nan
>>> np.array_equal(a, b, equal_nan=True)
True

log = <ufunc 'log'>

log2 = <ufunc 'log2'>

static any(a, axis=None, out=None, keepdims=<no value>, *, where=<no value>)

Test whether any array element along a given axis evaluates to True.

Returns single boolean if axis is None

Parameters:

• a (array_like) – Input array or object that can be converted to an array.

• axis (None or int or tuple of ints, optional) – Axis or axes along which a logical OR reduction is performed. The default (axis=None) is to perform a logical OR over all the dimensions of the input array. axis may be negative, in which case it counts from the last to the first axis.

  New in version 1.7.0.

  If this is a tuple of ints, a reduction is performed on multiple axes, instead of a single axis or all the axes as before.

• out (ndarray, optional) – Alternate output array in which to place the result. It must have the same shape as the expected output and its type is preserved (e.g., if it is of type float, then it will remain so, returning 1.0 for True and 0.0 for False, regardless of the type of a). See Output type determination for more details.

• keepdims (bool, optional) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.

  If the default value is passed, then keepdims will not be passed through to the any method of sub-classes of ndarray, however any non-default value will be. If the sub-class’ method does not implement keepdims any exceptions will be raised.

• where (array_like of bool, optional) – Elements to include in checking for any True values. See ~numpy.ufunc.reduce for details.

  New in version 1.20.0.

Returns:

any – A new boolean or ndarray is returned unless out is specified, in which case a reference to out is returned.

Return type:

bool or ndarray

See also

ndarray.any

equivalent method

all

Test whether all elements along a given axis evaluate to True.

Notes

Not a Number (NaN), positive infinity and negative infinity evaluate to True because these are not equal to zero.

Examples

>>> np.any([[True, False], [True, True]])
True

>>> np.any([[True, False], [False, False]], axis=0)
array([ True, False])

>>> np.any([-1, 0, 5])
True

>>> np.any(np.nan)
True

>>> np.any([[True, False], [False, False]], where=[[False], [True]])
False

>>> o=np.array(False)
>>> z=np.any([-1, 4, 5], out=o)
>>> z, o
(array(True), array(True))
>>> # Check now that z is a reference to o
>>> z is o
True
>>> id(z), id(o) # identity of z and o              
(191614240, 191614240)

static all(a, axis=None, out=None, keepdims=<no value>, *, where=<no value>)

Test whether all array elements along a given axis evaluate to True.

Parameters:

• a (array_like) – Input array or object that can be converted to an array.

• axis (None or int or tuple of ints, optional) – Axis or axes along which a logical AND reduction is performed. The default (axis=None) is to perform a logical AND over all the dimensions of the input array. axis may be negative, in which case it counts from the last to the first axis.

  New in version 1.7.0.

  If this is a tuple of ints, a reduction is performed on multiple axes, instead of a single axis or all the axes as before.

• out (ndarray, optional) – Alternate output array in which to place the result. It must have the same shape as the expected output and its type is preserved (e.g., if dtype(out) is float, the result will consist of 0.0’s and 1.0’s). See Output type determination for more details.

• keepdims (bool, optional) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.

  If the default value is passed, then keepdims will not be passed through to the all method of sub-classes of ndarray, however any non-default value will be. If the sub-class’ method does not implement keepdims any exceptions will be raised.

• where (array_like of bool, optional) – Elements to include in checking for all True values. See ~numpy.ufunc.reduce for details.

  New in version 1.20.0.

Returns:

all – A new boolean or array is returned unless out is specified, in which case a reference to out is returned.

Return type:

ndarray, bool

See also

ndarray.all

equivalent method

any

Test whether any element along a given axis evaluates to True.

Notes

Not a Number (NaN), positive infinity and negative infinity evaluate to True because these are not equal to zero.

Examples

>>> np.all([[True,False],[True,True]])
False

>>> np.all([[True,False],[True,True]], axis=0)
array([ True, False])

>>> np.all([-1, 4, 5])
True

>>> np.all([1.0, np.nan])
True

>>> np.all([[True, True], [False, True]], where=[[True], [False]])
True

>>> o=np.array(False)
>>> z=np.all([-1, 4, 5], out=o)
>>> id(z), id(o), z
(28293632, 28293632, array(True)) # may vary

static copy(a, order='K', subok=False)

Return an array copy of the given object.

Parameters:

• a (array_like) – Input data.

• order ({‘C’, ‘F’, ‘A’, ‘K’}, optional) – Controls the memory layout of the copy. ‘C’ means C-order, ‘F’ means F-order, ‘A’ means ‘F’ if a is Fortran contiguous, ‘C’ otherwise. ‘K’ means match the layout of a as closely as possible. (Note that this function and ndarray.copy() are very similar, but have different default values for their order= arguments.)

• subok (bool, optional) – If True, then sub-classes will be passed-through, otherwise the returned array will be forced to be a base-class array (defaults to False).

  New in version 1.19.0.

Returns:

arr – Array interpretation of a.

Return type:

ndarray

See also

ndarray.copy

Preferred method for creating an array copy

Notes

This is equivalent to:

>>> np.array(a, copy=True)  

Examples

Create an array x, with a reference y and a copy z:

>>> x = np.array([1, 2, 3])
>>> y = x
>>> z = np.copy(x)

Note that, when we modify x, y changes, but not z:

>>> x[0] = 10
>>> x[0] == y[0]
True
>>> x[0] == z[0]
False

Note that, np.copy clears previously set WRITEABLE=False flag.

>>> a = np.array([1, 2, 3])
>>> a.flags["WRITEABLE"] = False
>>> b = np.copy(a)
>>> b.flags["WRITEABLE"]
True
>>> b[0] = 3
>>> b
array([3, 2, 3])

Note that np.copy is a shallow copy and will not copy object elements within arrays. This is mainly important for arrays containing Python objects. The new array will contain the same object which may lead to surprises if that object can be modified (is mutable):

>>> a = np.array([1, 'm', [2, 3, 4]], dtype=object)
>>> b = np.copy(a)
>>> b[2][0] = 10
>>> a
array([1, 'm', list([10, 3, 4])], dtype=object)

To ensure all elements within an object array are copied, use copy.deepcopy:

>>> import copy
>>> a = np.array([1, 'm', [2, 3, 4]], dtype=object)
>>> c = copy.deepcopy(a)
>>> c[2][0] = 10
>>> c
array([1, 'm', list([10, 3, 4])], dtype=object)
>>> a
array([1, 'm', list([2, 3, 4])], dtype=object)

static nonzero(a)

Return the indices of the elements that are non-zero.

Returns a tuple of arrays, one for each dimension of a, containing the indices of the non-zero elements in that dimension. The values in a are always tested and returned in row-major, C-style order.

To group the indices by element, rather than dimension, use argwhere, which returns a row for each non-zero element.

Note

When called on a zero-d array or scalar, nonzero(a) is treated as nonzero(atleast_1d(a)).

Deprecated since version 1.17.0: Use atleast_1d explicitly if this behavior is deliberate.

Parameters:

a (array_like) – Input array.

Returns:

tuple_of_arrays – Indices of elements that are non-zero.

Return type:

tuple

See also

flatnonzero

Return indices that are non-zero in the flattened version of the input array.

ndarray.nonzero

Equivalent ndarray method.

count_nonzero

Counts the number of non-zero elements in the input array.

Notes

While the nonzero values can be obtained with a[nonzero(a)], it is recommended to use x[x.astype(bool)] or x[x != 0] instead, which will correctly handle 0-d arrays.

Examples

>>> x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]])
>>> x
array([[3, 0, 0],
       [0, 4, 0],
       [5, 6, 0]])
>>> np.nonzero(x)
(array([0, 1, 2, 2]), array([0, 1, 0, 1]))

>>> x[np.nonzero(x)]
array([3, 4, 5, 6])
>>> np.transpose(np.nonzero(x))
array([[0, 0],
       [1, 1],
       [2, 0],
       [2, 1]])

A common use for nonzero is to find the indices of an array, where a condition is True. Given an array a, the condition a > 3 is a boolean array and since False is interpreted as 0, np.nonzero(a > 3) yields the indices of the a where the condition is true.

>>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> a > 3
array([[False, False, False],
       [ True,  True,  True],
       [ True,  True,  True]])
>>> np.nonzero(a > 3)
(array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))

Using this result to index a is equivalent to using the mask directly:

>>> a[np.nonzero(a > 3)]
array([4, 5, 6, 7, 8, 9])
>>> a[a > 3]  # prefer this spelling
array([4, 5, 6, 7, 8, 9])

nonzero can also be called as a method of the array.

>>> (a > 3).nonzero()
(array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))

static astype(data, dtype, copy=True)

static tensor(data, device=NoParam)

static numpy(data)

static tolist(data)

static contiguous(data)

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

static asarray(data, dtype=None)

Cast data into a numpy representation

static ensure(data, dtype=None)

Cast data into a numpy representation

static dtype_kind(data)

static floor(data, out=None)

static ceil(data, out=None)

static ifloor(data, out=None)

static iceil(data, out=None)

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

static iround(data, out=None, dtype=<class 'int'>)

static clip(a, a_min, a_max, out=None, **kwargs)

Clip (limit) the values in an array.

Given an interval, values outside the interval are clipped to the interval edges. For example, if an interval of [0, 1] is specified, values smaller than 0 become 0, and values larger than 1 become 1.

Equivalent to but faster than np.minimum(a_max, np.maximum(a, a_min)).

No check is performed to ensure a_min < a_max.

Parameters:

• a (array_like) – Array containing elements to clip.

• a_min, a_max (array_like or None) – Minimum and maximum value. If None, clipping is not performed on the corresponding edge. Only one of a_min and a_max may be None. Both are broadcast against a.

• out (ndarray, optional) – The results will be placed in this array. It may be the input array for in-place clipping. out must be of the right shape to hold the output. Its type is preserved.

• **kwargs – For other keyword-only arguments, see the ufunc docs.

  New in version 1.17.0.

Returns:

clipped_array – An array with the elements of a, but where values < a_min are replaced with a_min, and those > a_max with a_max.

Return type:

ndarray

See also

Output type determination

Notes

When a_min is greater than a_max, clip returns an array in which all values are equal to a_max, as shown in the second example.

Examples

>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.clip(a, 1, 8)
array([1, 1, 2, 3, 4, 5, 6, 7, 8, 8])
>>> np.clip(a, 8, 1)
array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
>>> np.clip(a, 3, 6, out=a)
array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6])
>>> a
array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6])
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.clip(a, [3, 4, 1, 1, 1, 4, 4, 4, 4, 4], 8)
array([3, 4, 2, 3, 4, 5, 6, 7, 8, 8])

static softmax(data, axis=None)

static kron(a, b)

Outer product

Parameters:

• a (ndarray) – Array of shape (r0, r1, … rN)

• b (ndarray) – Array of shape (s0, s1, … sN)

Returns:

with shape (r0*s0, r1*s1, …, rN*SN).

Return type:

ndarray

Note

# From numpy doc kron(a,b)[k0,k1,…,kN] = a[i0,i1,…,iN] * b[j0,j1,…,jN]

kt = it * st + jt, t = 0,…,N

Notes

If a.shape = (r0,r1,..,rN) and b.shape = (s0,s1,…,sN), the Kronecker product has shape (r0*s0, r1*s1, …, rN*SN).

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)
>>> import kwarray
>>> import torch
>>> take = kwarray.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)
>>> import kwarray
>>> import torch
>>> compress = kwarray.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
>>> import torch
>>> 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

alias of TorchImpls

_numpy

alias of NumpyImpls

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).

static result_type(*arrays_and_dtypes)

Return type from promotion rules

Example

>>> import kwarray
>>> kwarray.ArrayAPI.result_type(float, np.uint8, np.float32, np.float16)
>>> # xdoctest: +REQUIRES(module:torch)
>>> import torch
>>> kwarray.ArrayAPI.result_type(torch.int32, torch.int64)

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

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

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

static take(data, *args, **kwargs)

static compress(data, *args, **kwargs)

static repeat(data, *args, **kwargs)

static tile(data, *args, **kwargs)

static view(data, *args, **kwargs)

static numel(data, *args, **kwargs)

static atleast_nd(data, *args, **kwargs)

static full_like(data, *args, **kwargs)

static ones_like(data, *args, **kwargs)

static zeros_like(data, *args, **kwargs)

static empty_like(data, *args, **kwargs)

static sum(data, *args, **kwargs)

static argsort(data, *args, **kwargs)

static argmax(data, *args, **kwargs)

static argmin(data, *args, **kwargs)

static max(data, *args, **kwargs)

static min(data, *args, **kwargs)

static max_argmax(data, *args, **kwargs)

static min_argmin(data, *args, **kwargs)

static maximum(data, *args, **kwargs)

static minimum(data, *args, **kwargs)

static matmul(data, *args, **kwargs)

static astype(data, *args, **kwargs)

static nonzero(data, *args, **kwargs)

static nan_to_num(data, *args, **kwargs)

static tensor(data, *args, **kwargs)

static numpy(data, *args, **kwargs)

static tolist(data, *args, **kwargs)

static asarray(data, *args, **kwargs)

static T(data, *args, **kwargs)

static transpose(data, *args, **kwargs)

static contiguous(data, *args, **kwargs)

static pad(data, *args, **kwargs)

static dtype_kind(data, *args, **kwargs)

static any(data, *args, **kwargs)

static all(data, *args, **kwargs)

static array_equal(data, *args, **kwargs)

static log2(data, *args, **kwargs)

static log(data, *args, **kwargs)

static copy(data, *args, **kwargs)

static iceil(data, *args, **kwargs)

static ifloor(data, *args, **kwargs)

static floor(data, *args, **kwargs)

static ceil(data, *args, **kwargs)

static round(data, *args, **kwargs)

static iround(data, *args, **kwargs)

static clip(data, *args, **kwargs)

static softmax(data, *args, **kwargs)

kwarray.arrayapi.TorchNumpyCompat

alias of ArrayAPI

kwarray.arrayapi._torch_dtype_lut()

kwarray.arrayapi.dtype_info(dtype)

Lookup datatype information

Parameters:

dtype (type) – a numpy, torch, or python numeric data type

Returns:

an iinfo of finfo structure depending on the input type.

Return type:

numpy.iinfo | numpy.finfo | torch.iinfo | torch.finfo

References

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

Example

>>> from kwarray.arrayapi import *  # NOQA
>>> try:
>>>     import torch
>>> except ImportError:
>>>     torch = None
>>> 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))