:py:mod:`kwarray.arrayapi` ========================== .. py:module:: kwarray.arrayapi .. autoapi-nested-parse:: 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: .. code:: python 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.(data)`` But you can also do ``kwarray.ArrayAPI.(data)`` on anything and it will do type checking and then do the operation you want. .. rubric:: 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) .. rubric:: Example >>> # xdoctest: +REQUIRES(module:torch) >>> import torch >>> import numpy as np >>> data1 = torch.rand(10, 10) >>> data2 = data1.numpy() Module Contents --------------- Classes ~~~~~~~ .. autoapisummary:: kwarray.arrayapi._ImplRegistry kwarray.arrayapi.TorchImpls kwarray.arrayapi.NumpyImpls kwarray.arrayapi.ArrayAPI Functions ~~~~~~~~~ .. autoapisummary:: kwarray.arrayapi._get_funcname kwarray.arrayapi._torch_dtype_lut kwarray.arrayapi.dtype_info Attributes ~~~~~~~~~~ .. autoapisummary:: kwarray.arrayapi.torch kwarray.arrayapi._REGISTERY kwarray.arrayapi._torchmethod kwarray.arrayapi._numpymethod kwarray.arrayapi._apimethod kwarray.arrayapi.TorchNumpyCompat .. py:data:: torch .. py:function:: _get_funcname(func) .. py:class:: _ImplRegistry Bases: :py:obj:`object` .. py:method:: _register(self, func, func_type, impl) .. py:method:: _implmethod(self, func=None, func_type='data_func', impl=None) .. py:method:: _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. .. py:method:: _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. .. py:data:: _REGISTERY .. py:data:: _torchmethod .. py:data:: _numpymethod .. py:data:: _apimethod .. py:class:: TorchImpls Bases: :py:obj:`object` Torch backend for the ArrayAPI API .. py:attribute:: is_tensor :annotation: = True .. py:attribute:: is_numpy :annotation: = False .. py:attribute:: ensure .. py:method:: cat(datas, axis=-1) .. py:method:: hstack(datas) Concatenates along axis=0 if inputs are are 1-D otherwise axis=1 .. rubric:: 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) .. py:method:: vstack(datas) Ensures that inputs datas are at least 2D (prepending a dimension of 1) and then concats along axis=0. .. rubric:: 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) .. py:method:: atleast_nd(arr, n, front=False) .. py:method:: view(data, *shape) .. py:method:: take(data, indices, axis=None) .. py:method:: compress(data, flags, axis=None) .. rubric:: 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,) .. py:method:: tile(data, reps) Implement np.tile in torch .. rubric:: 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) .. py:method:: repeat(data, repeats, axis=None) :abstractmethod: I'm not actually sure how to implement this efficiently .. rubric:: 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]) .. py:method:: T(data) .. py:method:: transpose(data, axes) .. rubric:: 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) .. py:method:: numel(data) .. py:method:: full_like(data, fill_value, dtype=None) .. py:method:: empty_like(data, dtype=None) .. py:method:: zeros_like(data, dtype=None) .. py:method:: ones_like(data, dtype=None) .. py:method:: full(shape, fill_value, dtype=float) .. py:method:: empty(shape, dtype=float) .. py:method:: zeros(shape, dtype=float) .. py:method:: ones(shape, dtype=float) .. py:method:: argmax(data, axis=None) .. py:method:: argsort(data, axis=-1, descending=False) .. rubric:: 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) .. py:method:: max(data, axis=None) .. rubric:: 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) .. py:method:: 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) .. py:method:: maximum(data1, data2, out=None) .. rubric:: 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) .. py:method:: minimum(data1, data2, out=None) .. rubric:: 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) .. py:method:: sum(data, axis=None) .. py:method:: nan_to_num(x, copy=True) .. py:method:: copy(data) .. py:method:: nonzero(data) .. py:method:: astype(data, dtype, copy=True) .. py:method:: tensor(data, device=ub.NoParam) .. py:method:: numpy(data) .. py:method:: tolist(data) .. py:method:: contiguous(data) .. py:method:: pad(data, pad_width, mode='constant') .. py:method:: asarray(data, dtype=None) Cast data into a tensor representation .. rubric:: Example >>> data = np.empty((2, 0, 196, 196), dtype=np.float32) .. py:method:: dtype_kind(data) returns the numpy code for the data type kind .. py:method:: floor(data, out=None) .. py:method:: ceil(data, out=None) .. py:method:: ifloor(data, out=None) .. py:method:: iceil(data, out=None) .. py:method:: round(data, decimals=0, out=None) .. rubric:: 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()) .. py:method:: iround(data, out=None, dtype=int) .. py:method:: clip(data, a_min=None, a_max=None, out=None) .. py:method:: softmax(data, axis=None) .. py:class:: NumpyImpls Bases: :py:obj:`object` Numpy backend for the ArrayAPI API .. py:attribute:: is_tensor :annotation: = False .. py:attribute:: is_numpy :annotation: = True .. py:attribute:: hstack .. py:attribute:: vstack .. py:attribute:: matmul .. py:attribute:: nan_to_num .. py:attribute:: log .. py:attribute:: log2 .. py:attribute:: any .. py:attribute:: all .. py:attribute:: copy .. py:attribute:: nonzero .. py:attribute:: ensure .. py:attribute:: clip .. py:method:: cat(datas, axis=-1) .. py:method:: atleast_nd(arr, n, front=False) .. py:method:: view(data, *shape) .. py:method:: take(data, indices, axis=None) .. py:method:: compress(data, flags, axis=None) .. py:method:: repeat(data, repeats, axis=None) .. py:method:: tile(data, reps) .. py:method:: T(data) .. py:method:: transpose(data, axes) .. py:method:: numel(data) .. py:method:: empty_like(data, dtype=None) .. py:method:: full_like(data, fill_value, dtype=None) .. py:method:: zeros_like(data, dtype=None) .. py:method:: ones_like(data, dtype=None) .. py:method:: full(shape, fill_value, dtype=float) .. py:method:: empty(shape, dtype=float) .. py:method:: zeros(shape, dtype=float) .. py:method:: ones(shape, dtype=float) .. py:method:: argmax(data, axis=None) .. py:method:: argsort(data, axis=-1, descending=False) .. py:method:: max(data, axis=None) .. py:method:: max_argmax(data, axis=None) .. py:method:: sum(data, axis=None) .. py:method:: maximum(data1, data2, out=None) .. py:method:: minimum(data1, data2, out=None) .. py:method:: astype(data, dtype, copy=True) .. py:method:: tensor(data, device=ub.NoParam) .. py:method:: numpy(data) .. py:method:: tolist(data) .. py:method:: contiguous(data) .. py:method:: pad(data, pad_width, mode='constant') .. py:method:: asarray(data, dtype=None) Cast data into a numpy representation .. py:method:: dtype_kind(data) .. py:method:: floor(data, out=None) .. py:method:: ceil(data, out=None) .. py:method:: ifloor(data, out=None) .. py:method:: iceil(data, out=None) .. py:method:: round(data, decimals=0, out=None) .. py:method:: iround(data, out=None, dtype=int) .. py:method:: softmax(data, axis=None) .. py:class:: ArrayAPI Bases: :py:obj:`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.`` 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 :func:`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. .. rubric:: 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)) .. rubric:: 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)) .. rubric:: 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]) .. py:attribute:: _torch .. py:attribute:: _numpy .. py:attribute:: take .. py:attribute:: compress .. py:attribute:: repeat .. py:attribute:: tile .. py:attribute:: view .. py:attribute:: numel .. py:attribute:: atleast_nd .. py:attribute:: full_like .. py:attribute:: ones_like .. py:attribute:: zeros_like .. py:attribute:: empty_like .. py:attribute:: sum .. py:attribute:: argmax .. py:attribute:: argsort .. py:attribute:: max .. py:attribute:: maximum .. py:attribute:: minimum .. py:attribute:: matmul .. py:attribute:: astype .. py:attribute:: nonzero .. py:attribute:: nan_to_num .. py:attribute:: tensor .. py:attribute:: numpy .. py:attribute:: tolist .. py:attribute:: asarray .. py:attribute:: asarray .. py:attribute:: T .. py:attribute:: transpose .. py:attribute:: contiguous .. py:attribute:: pad .. py:attribute:: dtype_kind .. py:attribute:: max_argmax .. py:attribute:: any .. py:attribute:: all .. py:attribute:: log2 .. py:attribute:: log .. py:attribute:: copy .. py:attribute:: iceil .. py:attribute:: ifloor .. py:attribute:: floor .. py:attribute:: ceil .. py:attribute:: round .. py:attribute:: iround .. py:attribute:: clip .. py:attribute:: softmax .. py:method:: impl(data) :staticmethod: Returns a namespace suitable for operating on the input data type :Parameters: **data** (*ndarray | Tensor*) -- data to be operated on .. py:method:: coerce(data) :staticmethod: Coerces some form of inputs into an array api (either numpy or torch). .. py:method:: cat(datas, *args, **kwargs) .. py:method:: hstack(datas, *args, **kwargs) .. py:method:: vstack(datas, *args, **kwargs) .. py:data:: TorchNumpyCompat .. py:function:: _torch_dtype_lut() .. py:function:: 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. :rtype: struct .. rubric:: References https://higra.readthedocs.io/en/stable/_modules/higra/hg_utils.html#dtype_info .. rubric:: 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))