Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,34 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from . import fft, linalg, random
|
||||
from ._dtypes import * # noqa: F403
|
||||
from ._funcs import * # noqa: F403
|
||||
from ._getlimits import finfo, iinfo
|
||||
from ._ndarray import (
|
||||
array,
|
||||
asarray,
|
||||
ascontiguousarray,
|
||||
can_cast,
|
||||
from_dlpack,
|
||||
ndarray,
|
||||
newaxis,
|
||||
result_type,
|
||||
)
|
||||
from ._ufuncs import * # noqa: F403
|
||||
from ._util import AxisError, UFuncTypeError
|
||||
|
||||
|
||||
from math import pi, e # usort: skip
|
||||
|
||||
|
||||
all = all # noqa: PLW0127
|
||||
alltrue = all
|
||||
|
||||
any = any # noqa: PLW0127
|
||||
sometrue = any
|
||||
|
||||
inf = float("inf")
|
||||
nan = float("nan")
|
||||
|
||||
False_ = False
|
||||
True_ = True
|
||||
@@ -0,0 +1,85 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Export torch work functions for binary ufuncs, rename/tweak to match numpy.
|
||||
This listing is further exported to public symbols in the `torch._numpy/_ufuncs.py` module.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import ( # noqa: F401
|
||||
add,
|
||||
arctan2,
|
||||
bitwise_and,
|
||||
bitwise_left_shift as left_shift,
|
||||
bitwise_or,
|
||||
bitwise_right_shift as right_shift,
|
||||
bitwise_xor,
|
||||
copysign,
|
||||
divide,
|
||||
eq as equal,
|
||||
float_power,
|
||||
floor_divide,
|
||||
fmax,
|
||||
fmin,
|
||||
fmod,
|
||||
gcd,
|
||||
greater,
|
||||
greater_equal,
|
||||
heaviside,
|
||||
hypot,
|
||||
lcm,
|
||||
ldexp,
|
||||
less,
|
||||
less_equal,
|
||||
logaddexp,
|
||||
logaddexp2,
|
||||
logical_and,
|
||||
logical_or,
|
||||
logical_xor,
|
||||
maximum,
|
||||
minimum,
|
||||
multiply,
|
||||
nextafter,
|
||||
not_equal,
|
||||
pow as power,
|
||||
remainder,
|
||||
remainder as mod,
|
||||
subtract,
|
||||
true_divide,
|
||||
)
|
||||
|
||||
from . import _dtypes_impl, _util
|
||||
|
||||
|
||||
# work around torch limitations w.r.t. numpy
|
||||
def matmul(x, y):
|
||||
# work around:
|
||||
# - RuntimeError: expected scalar type Int but found Double
|
||||
# - RuntimeError: "addmm_impl_cpu_" not implemented for 'Bool'
|
||||
# - RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
|
||||
dtype = _dtypes_impl.result_type_impl(x, y)
|
||||
is_bool = dtype == torch.bool
|
||||
is_half = (x.dtype == torch.float16 or y.dtype == torch.float16) and (
|
||||
x.is_cpu or y.is_cpu
|
||||
)
|
||||
|
||||
work_dtype = dtype
|
||||
if is_bool:
|
||||
work_dtype = torch.uint8
|
||||
if is_half:
|
||||
work_dtype = torch.float32
|
||||
|
||||
x = _util.cast_if_needed(x, work_dtype)
|
||||
y = _util.cast_if_needed(y, work_dtype)
|
||||
|
||||
result = torch.matmul(x, y)
|
||||
|
||||
if work_dtype != dtype:
|
||||
result = result.to(dtype)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# a stub implementation of divmod, should be improved after
|
||||
# https://github.com/pytorch/pytorch/issues/90820 is fixed in pytorch
|
||||
def divmod(x, y):
|
||||
return x // y, x % y
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,452 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Define analogs of numpy dtypes supported by pytorch.
|
||||
Define the scalar types and supported dtypes and numpy <--> torch dtype mappings.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl
|
||||
|
||||
|
||||
# ### Scalar types ###
|
||||
|
||||
|
||||
class generic:
|
||||
name = "generic"
|
||||
|
||||
def __new__(cls, value):
|
||||
# NumPy scalars are modelled as 0-D arrays
|
||||
# so a call to np.float32(4) produces a 0-D array.
|
||||
|
||||
from ._ndarray import asarray, ndarray
|
||||
|
||||
if isinstance(value, str) and value in ["inf", "nan"]:
|
||||
value = {"inf": torch.inf, "nan": torch.nan}[value]
|
||||
|
||||
if isinstance(value, ndarray):
|
||||
return value.astype(cls)
|
||||
else:
|
||||
return asarray(value, dtype=cls)
|
||||
|
||||
|
||||
##################
|
||||
# abstract types #
|
||||
##################
|
||||
|
||||
|
||||
class number(generic):
|
||||
name = "number"
|
||||
|
||||
|
||||
class integer(number):
|
||||
name = "integer"
|
||||
|
||||
|
||||
class inexact(number):
|
||||
name = "inexact"
|
||||
|
||||
|
||||
class signedinteger(integer):
|
||||
name = "signedinteger"
|
||||
|
||||
|
||||
class unsignedinteger(integer):
|
||||
name = "unsignedinteger"
|
||||
|
||||
|
||||
class floating(inexact):
|
||||
name = "floating"
|
||||
|
||||
|
||||
class complexfloating(inexact):
|
||||
name = "complexfloating"
|
||||
|
||||
|
||||
_abstract_dtypes = [
|
||||
"generic",
|
||||
"number",
|
||||
"integer",
|
||||
"signedinteger",
|
||||
"unsignedinteger",
|
||||
"inexact",
|
||||
"floating",
|
||||
"complexfloating",
|
||||
]
|
||||
|
||||
# ##### concrete types
|
||||
|
||||
# signed integers
|
||||
|
||||
|
||||
class int8(signedinteger):
|
||||
name = "int8"
|
||||
typecode = "b"
|
||||
torch_dtype = torch.int8
|
||||
|
||||
|
||||
class int16(signedinteger):
|
||||
name = "int16"
|
||||
typecode = "h"
|
||||
torch_dtype = torch.int16
|
||||
|
||||
|
||||
class int32(signedinteger):
|
||||
name = "int32"
|
||||
typecode = "i"
|
||||
torch_dtype = torch.int32
|
||||
|
||||
|
||||
class int64(signedinteger):
|
||||
name = "int64"
|
||||
typecode = "l"
|
||||
torch_dtype = torch.int64
|
||||
|
||||
|
||||
# unsigned integers
|
||||
|
||||
|
||||
class uint8(unsignedinteger):
|
||||
name = "uint8"
|
||||
typecode = "B"
|
||||
torch_dtype = torch.uint8
|
||||
|
||||
|
||||
class uint16(unsignedinteger):
|
||||
name = "uint16"
|
||||
typecode = "H"
|
||||
torch_dtype = torch.uint16
|
||||
|
||||
|
||||
class uint32(signedinteger):
|
||||
name = "uint32"
|
||||
typecode = "I"
|
||||
torch_dtype = torch.uint32
|
||||
|
||||
|
||||
class uint64(signedinteger):
|
||||
name = "uint64"
|
||||
typecode = "L"
|
||||
torch_dtype = torch.uint64
|
||||
|
||||
|
||||
# floating point
|
||||
|
||||
|
||||
class float16(floating):
|
||||
name = "float16"
|
||||
typecode = "e"
|
||||
torch_dtype = torch.float16
|
||||
|
||||
|
||||
class float32(floating):
|
||||
name = "float32"
|
||||
typecode = "f"
|
||||
torch_dtype = torch.float32
|
||||
|
||||
|
||||
class float64(floating):
|
||||
name = "float64"
|
||||
typecode = "d"
|
||||
torch_dtype = torch.float64
|
||||
|
||||
|
||||
class complex64(complexfloating):
|
||||
name = "complex64"
|
||||
typecode = "F"
|
||||
torch_dtype = torch.complex64
|
||||
|
||||
|
||||
class complex128(complexfloating):
|
||||
name = "complex128"
|
||||
typecode = "D"
|
||||
torch_dtype = torch.complex128
|
||||
|
||||
|
||||
class bool_(generic):
|
||||
name = "bool_"
|
||||
typecode = "?"
|
||||
torch_dtype = torch.bool
|
||||
|
||||
|
||||
# name aliases
|
||||
_name_aliases = {
|
||||
"intp": int64,
|
||||
"int_": int64,
|
||||
"intc": int32,
|
||||
"byte": int8,
|
||||
"short": int16,
|
||||
"longlong": int64, # XXX: is this correct?
|
||||
"ulonglong": uint64,
|
||||
"ubyte": uint8,
|
||||
"half": float16,
|
||||
"single": float32,
|
||||
"double": float64,
|
||||
"float_": float64,
|
||||
"csingle": complex64,
|
||||
"singlecomplex": complex64,
|
||||
"cdouble": complex128,
|
||||
"cfloat": complex128,
|
||||
"complex_": complex128,
|
||||
}
|
||||
# We register float_ = float32 and so on
|
||||
for name, obj in _name_aliases.items():
|
||||
vars()[name] = obj
|
||||
|
||||
|
||||
# Replicate this NumPy-defined way of grouping scalar types,
|
||||
# cf tests/core/test_scalar_methods.py
|
||||
sctypes = {
|
||||
"int": [int8, int16, int32, int64],
|
||||
"uint": [uint8, uint16, uint32, uint64],
|
||||
"float": [float16, float32, float64],
|
||||
"complex": [complex64, complex128],
|
||||
"others": [bool_],
|
||||
}
|
||||
|
||||
|
||||
# Support mappings/functions
|
||||
|
||||
_names = {st.name: st for cat in sctypes for st in sctypes[cat]}
|
||||
_typecodes = {st.typecode: st for cat in sctypes for st in sctypes[cat]}
|
||||
_torch_dtypes = {st.torch_dtype: st for cat in sctypes for st in sctypes[cat]}
|
||||
|
||||
|
||||
_aliases = {
|
||||
"u1": uint8,
|
||||
"i1": int8,
|
||||
"i2": int16,
|
||||
"i4": int32,
|
||||
"i8": int64,
|
||||
"b": int8, # XXX: srsly?
|
||||
"f2": float16,
|
||||
"f4": float32,
|
||||
"f8": float64,
|
||||
"c8": complex64,
|
||||
"c16": complex128,
|
||||
# numpy-specific trailing underscore
|
||||
"bool_": bool_,
|
||||
}
|
||||
|
||||
|
||||
_python_types = {
|
||||
int: int64,
|
||||
float: float64,
|
||||
complex: complex128,
|
||||
builtins.bool: bool_,
|
||||
# also allow stringified names of python types
|
||||
int.__name__: int64,
|
||||
float.__name__: float64,
|
||||
complex.__name__: complex128,
|
||||
builtins.bool.__name__: bool_,
|
||||
}
|
||||
|
||||
|
||||
def sctype_from_string(s):
|
||||
"""Normalize a string value: a type 'name' or a typecode or a width alias."""
|
||||
if s in _names:
|
||||
return _names[s]
|
||||
if s in _name_aliases:
|
||||
return _name_aliases[s]
|
||||
if s in _typecodes:
|
||||
return _typecodes[s]
|
||||
if s in _aliases:
|
||||
return _aliases[s]
|
||||
if s in _python_types:
|
||||
return _python_types[s]
|
||||
raise TypeError(f"data type {s!r} not understood")
|
||||
|
||||
|
||||
def sctype_from_torch_dtype(torch_dtype):
|
||||
return _torch_dtypes[torch_dtype]
|
||||
|
||||
|
||||
# ### DTypes. ###
|
||||
|
||||
|
||||
def dtype(arg):
|
||||
if arg is None:
|
||||
arg = _dtypes_impl.default_dtypes().float_dtype
|
||||
return DType(arg)
|
||||
|
||||
|
||||
class DType:
|
||||
def __init__(self, arg):
|
||||
# a pytorch object?
|
||||
if isinstance(arg, torch.dtype):
|
||||
sctype = _torch_dtypes[arg]
|
||||
elif isinstance(arg, torch.Tensor):
|
||||
sctype = _torch_dtypes[arg.dtype]
|
||||
# a scalar type?
|
||||
elif issubclass_(arg, generic):
|
||||
sctype = arg
|
||||
# a dtype already?
|
||||
elif isinstance(arg, DType):
|
||||
sctype = arg._scalar_type
|
||||
# a has a right attribute?
|
||||
elif hasattr(arg, "dtype"):
|
||||
sctype = arg.dtype._scalar_type
|
||||
else:
|
||||
sctype = sctype_from_string(arg)
|
||||
self._scalar_type = sctype
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._scalar_type.name
|
||||
|
||||
@property
|
||||
def type(self):
|
||||
return self._scalar_type
|
||||
|
||||
@property
|
||||
def kind(self):
|
||||
# https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
|
||||
return _torch_dtypes[self.torch_dtype].name[0]
|
||||
|
||||
@property
|
||||
def typecode(self):
|
||||
return self._scalar_type.typecode
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, DType):
|
||||
return self._scalar_type == other._scalar_type
|
||||
try:
|
||||
other_instance = DType(other)
|
||||
except TypeError:
|
||||
return False
|
||||
return self._scalar_type == other_instance._scalar_type
|
||||
|
||||
@property
|
||||
def torch_dtype(self):
|
||||
return self._scalar_type.torch_dtype
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self._scalar_type.name)
|
||||
|
||||
def __repr__(self):
|
||||
return f'dtype("{self.name}")'
|
||||
|
||||
__str__ = __repr__
|
||||
|
||||
@property
|
||||
def itemsize(self):
|
||||
elem = self.type(1)
|
||||
return elem.tensor.element_size()
|
||||
|
||||
def __getstate__(self):
|
||||
return self._scalar_type
|
||||
|
||||
def __setstate__(self, value):
|
||||
self._scalar_type = value
|
||||
|
||||
|
||||
typecodes = {
|
||||
"All": "efdFDBbhil?",
|
||||
"AllFloat": "efdFD",
|
||||
"AllInteger": "Bbhil",
|
||||
"Integer": "bhil",
|
||||
"UnsignedInteger": "B",
|
||||
"Float": "efd",
|
||||
"Complex": "FD",
|
||||
}
|
||||
|
||||
|
||||
# ### Defaults and dtype discovery
|
||||
|
||||
|
||||
def set_default_dtype(fp_dtype="numpy", int_dtype="numpy"):
|
||||
"""Set the (global) defaults for fp, complex, and int dtypes.
|
||||
|
||||
The complex dtype is inferred from the float (fp) dtype. It has
|
||||
a width at least twice the width of the float dtype,
|
||||
i.e., it's complex128 for float64 and complex64 for float32.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fp_dtype
|
||||
Allowed values are "numpy", "pytorch" or dtype_like things which
|
||||
can be converted into a DType instance.
|
||||
Default is "numpy" (i.e. float64).
|
||||
int_dtype
|
||||
Allowed values are "numpy", "pytorch" or dtype_like things which
|
||||
can be converted into a DType instance.
|
||||
Default is "numpy" (i.e. int64).
|
||||
|
||||
Returns
|
||||
-------
|
||||
The old default dtype state: a namedtuple with attributes ``float_dtype``,
|
||||
``complex_dtypes`` and ``int_dtype``. These attributes store *pytorch*
|
||||
dtypes.
|
||||
|
||||
Notes
|
||||
------------
|
||||
This functions has a side effect: it sets the global state with the provided dtypes.
|
||||
|
||||
The complex dtype has bit width of at least twice the width of the float
|
||||
dtype, i.e. it's complex128 for float64 and complex64 for float32.
|
||||
|
||||
"""
|
||||
if fp_dtype not in ["numpy", "pytorch"]:
|
||||
fp_dtype = dtype(fp_dtype).torch_dtype
|
||||
if int_dtype not in ["numpy", "pytorch"]:
|
||||
int_dtype = dtype(int_dtype).torch_dtype
|
||||
|
||||
if fp_dtype == "numpy":
|
||||
float_dtype = torch.float64
|
||||
elif fp_dtype == "pytorch":
|
||||
float_dtype = torch.float32
|
||||
else:
|
||||
float_dtype = fp_dtype
|
||||
|
||||
complex_dtype = {
|
||||
torch.float64: torch.complex128,
|
||||
torch.float32: torch.complex64,
|
||||
torch.float16: torch.complex64,
|
||||
}[float_dtype]
|
||||
|
||||
if int_dtype in ["numpy", "pytorch"]:
|
||||
int_dtype = torch.int64
|
||||
|
||||
new_defaults = _dtypes_impl.DefaultDTypes(
|
||||
float_dtype=float_dtype, complex_dtype=complex_dtype, int_dtype=int_dtype
|
||||
)
|
||||
|
||||
# set the new global state and return the old state
|
||||
old_defaults = _dtypes_impl.default_dtypes
|
||||
_dtypes_impl._default_dtypes = new_defaults
|
||||
return old_defaults
|
||||
|
||||
|
||||
def issubclass_(arg, klass):
|
||||
try:
|
||||
return issubclass(arg, klass)
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
|
||||
def issubdtype(arg1, arg2):
|
||||
# cf https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numerictypes.py#L356-L420
|
||||
|
||||
# We also accept strings even if NumPy doesn't as dtypes are serialized as their
|
||||
# string representation in dynamo's graph
|
||||
def str_to_abstract(t):
|
||||
if isinstance(t, str) and t in _abstract_dtypes:
|
||||
return globals()[t]
|
||||
return t
|
||||
|
||||
arg1 = str_to_abstract(arg1)
|
||||
arg2 = str_to_abstract(arg2)
|
||||
|
||||
if not issubclass_(arg1, generic):
|
||||
arg1 = dtype(arg1).type
|
||||
if not issubclass_(arg2, generic):
|
||||
arg2 = dtype(arg2).type
|
||||
return issubclass(arg1, arg2)
|
||||
|
||||
|
||||
__all__ = ["dtype", "DType", "typecodes", "issubdtype", "set_default_dtype", "sctypes"]
|
||||
__all__ += list(_names.keys()) # noqa: PLE0605
|
||||
__all__ += list(_name_aliases.keys()) # noqa: PLE0605
|
||||
__all__ += _abstract_dtypes # noqa: PLE0605
|
||||
@@ -0,0 +1,230 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Dtypes/scalar type implementations with torch dtypes.
|
||||
|
||||
Here `dtype` is always a torch.dtype, this module knows nothing about
|
||||
scalar types, wrapper dtypes or anything like that. PyTorch only.
|
||||
"""
|
||||
|
||||
from collections import namedtuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# defaults : mimic NumPy, allow user control
|
||||
DefaultDTypes = namedtuple(
|
||||
"DefaultDTypes", ["float_dtype", "complex_dtype", "int_dtype"]
|
||||
)
|
||||
|
||||
# a global state
|
||||
# We set it the first time we call default_dtypes() to avoid importing
|
||||
# torch._dynamo.config and create a circular reference
|
||||
_default_dtypes = None
|
||||
|
||||
|
||||
def default_dtypes():
|
||||
global _default_dtypes
|
||||
if _default_dtypes is None:
|
||||
import torch._dynamo.config as config
|
||||
|
||||
_default_dtypes = DefaultDTypes(
|
||||
float_dtype=getattr(torch, config.numpy_default_float),
|
||||
complex_dtype=getattr(torch, config.numpy_default_complex),
|
||||
int_dtype=getattr(torch, config.numpy_default_int),
|
||||
)
|
||||
if not isinstance(_default_dtypes.float_dtype, torch.dtype):
|
||||
raise AssertionError(
|
||||
f"float_dtype must be a torch.dtype, got {type(_default_dtypes.float_dtype)}"
|
||||
)
|
||||
if not isinstance(_default_dtypes.complex_dtype, torch.dtype):
|
||||
raise AssertionError(
|
||||
f"complex_dtype must be a torch.dtype, got {type(_default_dtypes.complex_dtype)}"
|
||||
)
|
||||
if not isinstance(_default_dtypes.int_dtype, torch.dtype):
|
||||
raise AssertionError(
|
||||
f"int_dtype must be a torch.dtype, got {type(_default_dtypes.int_dtype)}"
|
||||
)
|
||||
return _default_dtypes
|
||||
|
||||
|
||||
def get_default_dtype_for(dtype):
|
||||
"""Default scalar type given sctype category."""
|
||||
if dtype == torch.bool:
|
||||
return dtype
|
||||
if dtype.is_complex:
|
||||
return default_dtypes().complex_dtype
|
||||
if dtype.is_floating_point:
|
||||
return default_dtypes().float_dtype
|
||||
# else, it must be (some) integer
|
||||
return default_dtypes().int_dtype
|
||||
|
||||
|
||||
from . import _casting_dicts as _cd
|
||||
|
||||
|
||||
def can_cast_impl(from_torch_dtype, to_torch_dtype, casting):
|
||||
return _cd._can_cast_dict[casting][from_torch_dtype][to_torch_dtype]
|
||||
|
||||
|
||||
def result_type_impl(*tensors):
|
||||
# NB: torch dtypes here
|
||||
dtyp = tensors[0].dtype
|
||||
if len(tensors) == 1:
|
||||
return dtyp
|
||||
|
||||
for curr in tensors[1:]:
|
||||
dtyp = _cd._result_type_dict[dtyp][curr.dtype]
|
||||
|
||||
return dtyp
|
||||
|
||||
|
||||
def python_type_for_torch(dtyp):
|
||||
"""Get a python scalar type a torch dtype"""
|
||||
if dtyp.is_floating_point:
|
||||
typ = float
|
||||
elif dtyp.is_complex:
|
||||
typ = complex
|
||||
elif dtyp == torch.bool:
|
||||
typ = bool
|
||||
else:
|
||||
typ = int
|
||||
return typ
|
||||
|
||||
|
||||
# ### NEP 50 helpers ###
|
||||
|
||||
_SCALAR_TYPES = (int, bool, float, complex)
|
||||
|
||||
_SCALAR_AND_SYMBOLIC_TYPES = (
|
||||
*_SCALAR_TYPES,
|
||||
torch.SymInt,
|
||||
torch.SymFloat,
|
||||
torch.SymBool,
|
||||
)
|
||||
|
||||
_NEP50_FUNCS_TENSOR_ONLY = (
|
||||
"minimum",
|
||||
"maximum",
|
||||
"logaddexp",
|
||||
"logaddexp2",
|
||||
"lcm",
|
||||
"gcd",
|
||||
"hypot",
|
||||
"heaviside",
|
||||
"fmod",
|
||||
"fmin",
|
||||
"fmax",
|
||||
"copysign",
|
||||
"arctan2",
|
||||
)
|
||||
|
||||
|
||||
def is_scalar(x):
|
||||
return isinstance(x, _SCALAR_TYPES)
|
||||
|
||||
|
||||
def is_scalar_or_symbolic(x):
|
||||
return isinstance(x, _SCALAR_AND_SYMBOLIC_TYPES)
|
||||
|
||||
|
||||
def _dtype_for_scalar(py_type):
|
||||
return {
|
||||
bool: torch.bool,
|
||||
torch.SymBool: torch.bool,
|
||||
int: torch.int64,
|
||||
torch.SymInt: torch.int64,
|
||||
float: torch.float64,
|
||||
torch.SymFloat: torch.float64,
|
||||
complex: torch.complex128,
|
||||
}[py_type]
|
||||
|
||||
|
||||
def _dtype_for_scalar_or_tensor(x):
|
||||
return x.dtype if isinstance(x, torch.Tensor) else _dtype_for_scalar(type(x))
|
||||
|
||||
|
||||
def is_float_or_fp_tensor(x):
|
||||
return _dtype_for_scalar_or_tensor(x).is_floating_point
|
||||
|
||||
|
||||
def is_complex_or_complex_tensor(x):
|
||||
return _dtype_for_scalar_or_tensor(x).is_complex
|
||||
|
||||
|
||||
def _category(dtype):
|
||||
return {
|
||||
torch.bool: 0,
|
||||
torch.SymBool: 0,
|
||||
# int
|
||||
torch.uint8: 1,
|
||||
torch.int8: 1,
|
||||
torch.int16: 1,
|
||||
torch.int32: 1,
|
||||
torch.int64: 1,
|
||||
torch.SymInt: 1,
|
||||
# float
|
||||
torch.float16: 2,
|
||||
torch.float32: 2,
|
||||
torch.float64: 2,
|
||||
torch.SymFloat: 2,
|
||||
# complex
|
||||
torch.complex64: 3,
|
||||
torch.complex128: 3,
|
||||
}[dtype]
|
||||
|
||||
|
||||
def nep50_to_tensors(x1, x2, handle_weaks, function_name):
|
||||
"""If either of inputs is a python scalar, type-promote with NEP 50."""
|
||||
|
||||
def to_tensor(scalar, dtype=None):
|
||||
if dtype is None:
|
||||
dtype = _dtype_for_scalar(type(scalar))
|
||||
dtype = get_default_dtype_for(dtype)
|
||||
return torch.as_tensor(scalar, dtype=dtype)
|
||||
|
||||
x1_is_weak = not isinstance(x1, torch.Tensor)
|
||||
x2_is_weak = not isinstance(x2, torch.Tensor)
|
||||
if not handle_weaks or (x1_is_weak and x2_is_weak):
|
||||
x1 = to_tensor(x1) if x1_is_weak else x1
|
||||
x2 = to_tensor(x2) if x2_is_weak else x2
|
||||
return x1, x2
|
||||
|
||||
# scalar <op> tensor: NEP 50
|
||||
if x1_is_weak == x2_is_weak:
|
||||
raise AssertionError(
|
||||
f"Expected exactly one weak type, got x1_is_weak={x1_is_weak}, x2_is_weak={x2_is_weak}"
|
||||
)
|
||||
|
||||
weak, not_weak = (x1, x2) if x1_is_weak else (x2, x1)
|
||||
|
||||
# find the dtype for the weak's type
|
||||
weak_dtype = _dtype_for_scalar(type(weak))
|
||||
|
||||
cat_weak = _category(weak_dtype)
|
||||
cat_not_weak = _category(not_weak.dtype)
|
||||
|
||||
dt = not_weak.dtype if cat_weak <= cat_not_weak else None
|
||||
|
||||
# special-case complex + float32
|
||||
if weak_dtype.is_complex and not_weak.dtype == torch.float32:
|
||||
dt = torch.complex64
|
||||
|
||||
# detect overflows: in PyTorch, uint8(-1) wraps around to 255,
|
||||
# while NEP50 mandates an exception.
|
||||
#
|
||||
# Note that we only check if each element of the binop overflows,
|
||||
# not the result. Consider, e.g. `uint8(100) + 200`. Operands are OK
|
||||
# in uint8, but the result overflows and wrap around 255.
|
||||
# Numpy emits a RuntimeWarning, PyTorch does not, and we do not either.
|
||||
if cat_weak == 1 and cat_not_weak == 1:
|
||||
# integers
|
||||
iinfo = torch.iinfo(not_weak.dtype)
|
||||
if not (iinfo.min <= weak <= iinfo.max):
|
||||
raise OverflowError(
|
||||
f"Python integer {weak} out of bounds for {not_weak.dtype}"
|
||||
)
|
||||
if weak_dtype != dt or function_name in _NEP50_FUNCS_TENSOR_ONLY:
|
||||
# finally, can make `weak` into a 0D tensor, if both parameters are required to be tensor.
|
||||
weak = to_tensor(weak, dt)
|
||||
|
||||
return (weak, not_weak) if x1_is_weak else (not_weak, weak)
|
||||
@@ -0,0 +1,76 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import inspect
|
||||
import itertools
|
||||
|
||||
from . import _funcs_impl, _reductions_impl
|
||||
from ._normalizations import normalizer
|
||||
|
||||
|
||||
# _funcs_impl.py contains functions which mimic NumPy's eponymous equivalents,
|
||||
# and consume/return PyTorch tensors/dtypes.
|
||||
# They are also type annotated.
|
||||
# Pull these functions from _funcs_impl and decorate them with @normalizer, which
|
||||
# - Converts any input `np.ndarray`, `torch._numpy.ndarray`, list of lists, Python scalars, etc into a `torch.Tensor`.
|
||||
# - Maps NumPy dtypes to PyTorch dtypes
|
||||
# - If the input to the `axis` kwarg is an ndarray, it maps it into a tuple
|
||||
# - Implements the semantics for the `out=` arg
|
||||
# - Wraps back the outputs into `torch._numpy.ndarrays`
|
||||
|
||||
|
||||
def _public_functions(mod):
|
||||
def is_public_function(f):
|
||||
return inspect.isfunction(f) and not f.__name__.startswith("_")
|
||||
|
||||
return inspect.getmembers(mod, is_public_function)
|
||||
|
||||
|
||||
# We fill in __all__ in the loop below
|
||||
__all__ = []
|
||||
|
||||
# decorate implementer functions with argument normalizers and export to the top namespace
|
||||
for name, func in itertools.chain(
|
||||
_public_functions(_funcs_impl), _public_functions(_reductions_impl)
|
||||
):
|
||||
if name in ["percentile", "quantile", "median"]:
|
||||
decorated = normalizer(func, promote_scalar_result=True)
|
||||
elif name == "einsum":
|
||||
# normalized manually
|
||||
decorated = func
|
||||
else:
|
||||
decorated = normalizer(func)
|
||||
|
||||
decorated.__qualname__ = name
|
||||
decorated.__name__ = name
|
||||
vars()[name] = decorated
|
||||
__all__.append(name)
|
||||
|
||||
|
||||
"""
|
||||
Vendored objects from numpy.lib.index_tricks
|
||||
"""
|
||||
|
||||
|
||||
class IndexExpression:
|
||||
"""
|
||||
Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
|
||||
last revision: 1999-7-23
|
||||
|
||||
Cosmetic changes by T. Oliphant 2001
|
||||
"""
|
||||
|
||||
def __init__(self, maketuple):
|
||||
self.maketuple = maketuple
|
||||
|
||||
def __getitem__(self, item):
|
||||
if self.maketuple and not isinstance(item, tuple):
|
||||
return (item,)
|
||||
else:
|
||||
return item
|
||||
|
||||
|
||||
index_exp = IndexExpression(maketuple=True)
|
||||
s_ = IndexExpression(maketuple=False)
|
||||
|
||||
|
||||
__all__ += ["index_exp", "s_"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,15 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes
|
||||
|
||||
|
||||
def finfo(dtyp):
|
||||
torch_dtype = _dtypes.dtype(dtyp).torch_dtype
|
||||
return torch.finfo(torch_dtype)
|
||||
|
||||
|
||||
def iinfo(dtyp):
|
||||
torch_dtype = _dtypes.dtype(dtyp).torch_dtype
|
||||
return torch.iinfo(torch_dtype)
|
||||
@@ -0,0 +1,733 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import math
|
||||
import operator
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes, _dtypes_impl, _funcs, _ufuncs, _util
|
||||
from ._normalizations import (
|
||||
ArrayLike,
|
||||
normalize_array_like,
|
||||
normalizer,
|
||||
NotImplementedType,
|
||||
)
|
||||
|
||||
|
||||
newaxis = None
|
||||
|
||||
FLAGS = [
|
||||
"C_CONTIGUOUS",
|
||||
"F_CONTIGUOUS",
|
||||
"OWNDATA",
|
||||
"WRITEABLE",
|
||||
"ALIGNED",
|
||||
"WRITEBACKIFCOPY",
|
||||
"FNC",
|
||||
"FORC",
|
||||
"BEHAVED",
|
||||
"CARRAY",
|
||||
"FARRAY",
|
||||
]
|
||||
|
||||
SHORTHAND_TO_FLAGS = {
|
||||
"C": "C_CONTIGUOUS",
|
||||
"F": "F_CONTIGUOUS",
|
||||
"O": "OWNDATA",
|
||||
"W": "WRITEABLE",
|
||||
"A": "ALIGNED",
|
||||
"X": "WRITEBACKIFCOPY",
|
||||
"B": "BEHAVED",
|
||||
"CA": "CARRAY",
|
||||
"FA": "FARRAY",
|
||||
}
|
||||
|
||||
|
||||
class Flags:
|
||||
def __init__(self, flag_to_value: dict):
|
||||
invalid_keys = [k for k in flag_to_value if k not in FLAGS]
|
||||
if invalid_keys:
|
||||
raise AssertionError(f"Invalid flag keys: {invalid_keys}")
|
||||
self._flag_to_value = flag_to_value
|
||||
|
||||
def __getattr__(self, attr: str):
|
||||
if attr.islower() and attr.upper() in FLAGS:
|
||||
return self[attr.upper()]
|
||||
else:
|
||||
raise AttributeError(f"No flag attribute '{attr}'")
|
||||
|
||||
def __getitem__(self, key):
|
||||
if key in SHORTHAND_TO_FLAGS:
|
||||
key = SHORTHAND_TO_FLAGS[key]
|
||||
if key in FLAGS:
|
||||
try:
|
||||
return self._flag_to_value[key]
|
||||
except KeyError as e:
|
||||
raise NotImplementedError(f"{key=}") from e
|
||||
else:
|
||||
raise KeyError(f"No flag key '{key}'")
|
||||
|
||||
def __setattr__(self, attr, value):
|
||||
if attr.islower() and attr.upper() in FLAGS:
|
||||
self[attr.upper()] = value
|
||||
else:
|
||||
super().__setattr__(attr, value)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if key in FLAGS or key in SHORTHAND_TO_FLAGS:
|
||||
raise NotImplementedError("Modifying flags is not implemented")
|
||||
else:
|
||||
raise KeyError(f"No flag key '{key}'")
|
||||
|
||||
|
||||
def create_method(fn, name=None):
|
||||
name = name or fn.__name__
|
||||
|
||||
def f(*args, **kwargs):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
f.__name__ = name
|
||||
f.__qualname__ = f"ndarray.{name}"
|
||||
return f
|
||||
|
||||
|
||||
# Map ndarray.name_method -> np.name_func
|
||||
# If name_func == None, it means that name_method == name_func
|
||||
methods = {
|
||||
"clip": None,
|
||||
"nonzero": None,
|
||||
"repeat": None,
|
||||
"round": None,
|
||||
"squeeze": None,
|
||||
"swapaxes": None,
|
||||
"ravel": None,
|
||||
# linalg
|
||||
"diagonal": None,
|
||||
"dot": None,
|
||||
"trace": None,
|
||||
# sorting
|
||||
"argsort": None,
|
||||
"searchsorted": None,
|
||||
# reductions
|
||||
"argmax": None,
|
||||
"argmin": None,
|
||||
"any": None,
|
||||
"all": None,
|
||||
"max": None,
|
||||
"min": None,
|
||||
"ptp": None,
|
||||
"sum": None,
|
||||
"prod": None,
|
||||
"mean": None,
|
||||
"var": None,
|
||||
"std": None,
|
||||
# scans
|
||||
"cumsum": None,
|
||||
"cumprod": None,
|
||||
# advanced indexing
|
||||
"take": None,
|
||||
"choose": None,
|
||||
}
|
||||
|
||||
dunder = {
|
||||
"abs": "absolute",
|
||||
"invert": None,
|
||||
"pos": "positive",
|
||||
"neg": "negative",
|
||||
"gt": "greater",
|
||||
"lt": "less",
|
||||
"ge": "greater_equal",
|
||||
"le": "less_equal",
|
||||
}
|
||||
|
||||
# dunder methods with right-looking and in-place variants
|
||||
ri_dunder = {
|
||||
"add": None,
|
||||
"sub": "subtract",
|
||||
"mul": "multiply",
|
||||
"truediv": "divide",
|
||||
"floordiv": "floor_divide",
|
||||
"pow": "power",
|
||||
"mod": "remainder",
|
||||
"and": "bitwise_and",
|
||||
"or": "bitwise_or",
|
||||
"xor": "bitwise_xor",
|
||||
"lshift": "left_shift",
|
||||
"rshift": "right_shift",
|
||||
"matmul": None,
|
||||
}
|
||||
|
||||
|
||||
def _upcast_int_indices(index):
|
||||
if isinstance(index, torch.Tensor):
|
||||
if index.dtype in (torch.int8, torch.int16, torch.int32, torch.uint8):
|
||||
return index.to(torch.int64)
|
||||
elif isinstance(index, tuple):
|
||||
return tuple(_upcast_int_indices(i) for i in index)
|
||||
return index
|
||||
|
||||
|
||||
def _has_advanced_indexing(index):
|
||||
"""Check if there's any advanced indexing"""
|
||||
return any(
|
||||
isinstance(idx, (Sequence, bool))
|
||||
or (isinstance(idx, torch.Tensor) and (idx.dtype == torch.bool or idx.ndim > 0))
|
||||
for idx in index
|
||||
)
|
||||
|
||||
|
||||
def _numpy_compatible_indexing(index):
|
||||
"""Convert scalar indices to lists when advanced indexing is present for NumPy compatibility."""
|
||||
if not isinstance(index, tuple):
|
||||
index = (index,)
|
||||
|
||||
# Check if there's any advanced indexing (sequences, booleans, or tensors)
|
||||
has_advanced = _has_advanced_indexing(index)
|
||||
|
||||
if not has_advanced:
|
||||
return index
|
||||
|
||||
# Convert integer scalar indices to single-element lists when advanced indexing is present
|
||||
# Note: Do NOT convert boolean scalars (True/False) as they have special meaning in NumPy
|
||||
converted = []
|
||||
for idx in index:
|
||||
if isinstance(idx, int) and not isinstance(idx, bool):
|
||||
# Integer scalars should be converted to lists
|
||||
converted.append([idx])
|
||||
elif (
|
||||
isinstance(idx, torch.Tensor)
|
||||
and idx.ndim == 0
|
||||
and not torch.is_floating_point(idx)
|
||||
and idx.dtype != torch.bool
|
||||
):
|
||||
# Zero-dimensional tensors holding integers should be treated the same as integer scalars
|
||||
converted.append([idx])
|
||||
else:
|
||||
# Everything else (booleans, lists, slices, etc.) stays as is
|
||||
converted.append(idx)
|
||||
|
||||
return tuple(converted)
|
||||
|
||||
|
||||
def _get_bool_depth(s):
|
||||
"""Returns the depth of a boolean sequence/tensor"""
|
||||
if isinstance(s, bool):
|
||||
return True, 0
|
||||
if isinstance(s, torch.Tensor) and s.dtype == torch.bool:
|
||||
return True, s.ndim
|
||||
if not (isinstance(s, Sequence) and s and s[0] != s):
|
||||
return False, 0
|
||||
is_bool, depth = _get_bool_depth(s[0])
|
||||
return is_bool, depth + 1
|
||||
|
||||
|
||||
def _numpy_empty_ellipsis_patch(index, tensor_ndim):
|
||||
"""
|
||||
Patch for NumPy-compatible ellipsis behavior when ellipsis doesn't match any dimensions.
|
||||
|
||||
In NumPy, when an ellipsis (...) doesn't actually match any dimensions of the input array,
|
||||
it still acts as a separator between advanced indices. PyTorch doesn't have this behavior.
|
||||
|
||||
This function detects when we have:
|
||||
1. Advanced indexing on both sides of an ellipsis
|
||||
2. The ellipsis doesn't actually match any dimensions
|
||||
"""
|
||||
if not isinstance(index, tuple):
|
||||
index = (index,)
|
||||
|
||||
# Find ellipsis position
|
||||
ellipsis_pos = None
|
||||
for i, idx in enumerate(index):
|
||||
if idx is Ellipsis:
|
||||
ellipsis_pos = i
|
||||
break
|
||||
|
||||
# If no ellipsis, no patch needed
|
||||
if ellipsis_pos is None:
|
||||
return index, lambda x: x, lambda x: x
|
||||
|
||||
# Count non-ellipsis dimensions consumed by the index
|
||||
consumed_dims = 0
|
||||
for idx in index:
|
||||
is_bool, depth = _get_bool_depth(idx)
|
||||
if is_bool:
|
||||
consumed_dims += depth
|
||||
elif idx is Ellipsis or idx is None:
|
||||
continue
|
||||
else:
|
||||
consumed_dims += 1
|
||||
|
||||
# Calculate how many dimensions the ellipsis should match
|
||||
ellipsis_dims = tensor_ndim - consumed_dims
|
||||
|
||||
# Check if ellipsis doesn't match any dimensions
|
||||
if ellipsis_dims == 0:
|
||||
# Check if we have advanced indexing on both sides of ellipsis
|
||||
left_advanced = _has_advanced_indexing(index[:ellipsis_pos])
|
||||
right_advanced = _has_advanced_indexing(index[ellipsis_pos + 1 :])
|
||||
|
||||
if left_advanced and right_advanced:
|
||||
# This is the case where NumPy and PyTorch differ
|
||||
# We need to ensure the advanced indices are treated as separated
|
||||
new_index = index[:ellipsis_pos] + (None,) + index[ellipsis_pos + 1 :]
|
||||
end_ndims = 1 + sum(
|
||||
1 for idx in index[ellipsis_pos + 1 :] if isinstance(idx, slice)
|
||||
)
|
||||
|
||||
def squeeze_fn(x):
|
||||
return x.squeeze(-end_ndims)
|
||||
|
||||
def unsqueeze_fn(x):
|
||||
if isinstance(x, torch.Tensor) and x.ndim >= end_ndims:
|
||||
return x.unsqueeze(-end_ndims)
|
||||
return x
|
||||
|
||||
return new_index, squeeze_fn, unsqueeze_fn
|
||||
|
||||
return index, lambda x: x, lambda x: x
|
||||
|
||||
|
||||
# Used to indicate that a parameter is unspecified (as opposed to explicitly
|
||||
# `None`)
|
||||
class _Unspecified:
|
||||
pass
|
||||
|
||||
|
||||
_Unspecified.unspecified = _Unspecified()
|
||||
|
||||
###############################################################
|
||||
# ndarray class #
|
||||
###############################################################
|
||||
|
||||
|
||||
class ndarray:
|
||||
def __init__(self, t=None):
|
||||
if t is None:
|
||||
self.tensor = torch.Tensor()
|
||||
elif isinstance(t, torch.Tensor):
|
||||
self.tensor = t
|
||||
else:
|
||||
raise ValueError(
|
||||
"ndarray constructor is not recommended; prefer"
|
||||
"either array(...) or zeros/empty(...)"
|
||||
)
|
||||
|
||||
# Register NumPy functions as methods
|
||||
for method, name in methods.items():
|
||||
fn = getattr(_funcs, name or method)
|
||||
vars()[method] = create_method(fn, method)
|
||||
|
||||
# Regular methods but coming from ufuncs
|
||||
conj = create_method(_ufuncs.conjugate, "conj")
|
||||
conjugate = create_method(_ufuncs.conjugate)
|
||||
|
||||
for method, name in dunder.items():
|
||||
fn = getattr(_ufuncs, name or method)
|
||||
method = f"__{method}__"
|
||||
vars()[method] = create_method(fn, method)
|
||||
|
||||
for method, name in ri_dunder.items():
|
||||
fn = getattr(_ufuncs, name or method)
|
||||
plain = f"__{method}__"
|
||||
vars()[plain] = create_method(fn, plain)
|
||||
rvar = f"__r{method}__"
|
||||
vars()[rvar] = create_method(lambda self, other, fn=fn: fn(other, self), rvar)
|
||||
ivar = f"__i{method}__"
|
||||
vars()[ivar] = create_method(
|
||||
lambda self, other, fn=fn: fn(self, other, out=self), ivar
|
||||
)
|
||||
|
||||
# There's no __idivmod__
|
||||
__divmod__ = create_method(_ufuncs.divmod, "__divmod__")
|
||||
__rdivmod__ = create_method(
|
||||
lambda self, other: _ufuncs.divmod(other, self), "__rdivmod__"
|
||||
)
|
||||
|
||||
# prevent loop variables leaking into the ndarray class namespace
|
||||
del ivar, rvar, name, plain, fn, method
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return tuple(self.tensor.shape)
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return self.tensor.numel()
|
||||
|
||||
@property
|
||||
def ndim(self):
|
||||
return self.tensor.ndim
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return _dtypes.dtype(self.tensor.dtype)
|
||||
|
||||
@property
|
||||
def strides(self):
|
||||
elsize = self.tensor.element_size()
|
||||
return tuple(stride * elsize for stride in self.tensor.stride())
|
||||
|
||||
@property
|
||||
def itemsize(self):
|
||||
return self.tensor.element_size()
|
||||
|
||||
@property
|
||||
def flags(self):
|
||||
# Note contiguous in torch is assumed C-style
|
||||
return Flags(
|
||||
{
|
||||
"C_CONTIGUOUS": self.tensor.is_contiguous(),
|
||||
"F_CONTIGUOUS": self.T.tensor.is_contiguous(),
|
||||
"OWNDATA": self.tensor._base is None,
|
||||
"WRITEABLE": True, # pytorch does not have readonly tensors
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
return self.tensor.data_ptr()
|
||||
|
||||
@property
|
||||
def nbytes(self):
|
||||
return self.tensor.storage().nbytes()
|
||||
|
||||
@property
|
||||
def T(self):
|
||||
return self.transpose()
|
||||
|
||||
@property
|
||||
def real(self):
|
||||
return _funcs.real(self)
|
||||
|
||||
@real.setter
|
||||
def real(self, value):
|
||||
self.tensor.real = asarray(value).tensor
|
||||
|
||||
@property
|
||||
def imag(self):
|
||||
return _funcs.imag(self)
|
||||
|
||||
@imag.setter
|
||||
def imag(self, value):
|
||||
self.tensor.imag = asarray(value).tensor
|
||||
|
||||
@property
|
||||
def flat(self):
|
||||
return self.ravel()
|
||||
|
||||
# ctors
|
||||
def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
|
||||
if order != "K":
|
||||
raise NotImplementedError(f"astype(..., order={order} is not implemented.")
|
||||
if casting != "unsafe":
|
||||
raise NotImplementedError(
|
||||
f"astype(..., casting={casting} is not implemented."
|
||||
)
|
||||
if not subok:
|
||||
raise NotImplementedError(f"astype(..., subok={subok} is not implemented.")
|
||||
if not copy:
|
||||
raise NotImplementedError(f"astype(..., copy={copy} is not implemented.")
|
||||
torch_dtype = _dtypes.dtype(dtype).torch_dtype
|
||||
t = self.tensor.to(torch_dtype)
|
||||
return ndarray(t)
|
||||
|
||||
@normalizer
|
||||
def copy(self: ArrayLike, order: NotImplementedType = "C"):
|
||||
return self.clone()
|
||||
|
||||
@normalizer
|
||||
def flatten(self: ArrayLike, order: NotImplementedType = "C"):
|
||||
return torch.flatten(self)
|
||||
|
||||
def resize(self, *new_shape, refcheck=False):
|
||||
# NB: differs from np.resize: fills with zeros instead of making repeated copies of input.
|
||||
if refcheck:
|
||||
raise NotImplementedError(
|
||||
f"resize(..., refcheck={refcheck} is not implemented."
|
||||
)
|
||||
if new_shape in [(), (None,)]:
|
||||
return
|
||||
|
||||
# support both x.resize((2, 2)) and x.resize(2, 2)
|
||||
if len(new_shape) == 1:
|
||||
new_shape = new_shape[0]
|
||||
if isinstance(new_shape, int):
|
||||
new_shape = (new_shape,)
|
||||
|
||||
if builtins.any(x < 0 for x in new_shape):
|
||||
raise ValueError("all elements of `new_shape` must be non-negative")
|
||||
|
||||
new_numel, old_numel = math.prod(new_shape), self.tensor.numel()
|
||||
|
||||
self.tensor.resize_(new_shape)
|
||||
|
||||
if new_numel >= old_numel:
|
||||
# zero-fill new elements
|
||||
if not self.tensor.is_contiguous():
|
||||
raise AssertionError("tensor must be contiguous for resize with growth")
|
||||
b = self.tensor.flatten() # does not copy
|
||||
b[old_numel:].zero_()
|
||||
|
||||
def view(self, dtype=_Unspecified.unspecified, type=_Unspecified.unspecified):
|
||||
if dtype is _Unspecified.unspecified:
|
||||
dtype = self.dtype
|
||||
if type is not _Unspecified.unspecified:
|
||||
raise NotImplementedError(f"view(..., type={type} is not implemented.")
|
||||
torch_dtype = _dtypes.dtype(dtype).torch_dtype
|
||||
tview = self.tensor.view(torch_dtype)
|
||||
return ndarray(tview)
|
||||
|
||||
@normalizer
|
||||
def fill(self, value: ArrayLike):
|
||||
# Both Pytorch and NumPy accept 0D arrays/tensors and scalars, and
|
||||
# error out on D > 0 arrays
|
||||
self.tensor.fill_(value)
|
||||
|
||||
def tolist(self):
|
||||
return self.tensor.tolist()
|
||||
|
||||
def __iter__(self):
|
||||
return (ndarray(x) for x in self.tensor.__iter__())
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
str(self.tensor)
|
||||
.replace("tensor", "torch.ndarray")
|
||||
.replace("dtype=torch.", "dtype=")
|
||||
)
|
||||
|
||||
__repr__ = create_method(__str__)
|
||||
|
||||
def __eq__(self, other):
|
||||
try:
|
||||
return _ufuncs.equal(self, other)
|
||||
except (RuntimeError, TypeError):
|
||||
# Failed to convert other to array: definitely not equal.
|
||||
falsy = torch.full(self.shape, fill_value=False, dtype=bool)
|
||||
return asarray(falsy)
|
||||
|
||||
def __ne__(self, other):
|
||||
return ~(self == other)
|
||||
|
||||
def __index__(self):
|
||||
try:
|
||||
return operator.index(self.tensor.item())
|
||||
except Exception as exc:
|
||||
raise TypeError(
|
||||
"only integer scalar arrays can be converted to a scalar index"
|
||||
) from exc
|
||||
|
||||
def __bool__(self):
|
||||
return bool(self.tensor)
|
||||
|
||||
def __int__(self):
|
||||
return int(self.tensor)
|
||||
|
||||
def __float__(self):
|
||||
return float(self.tensor)
|
||||
|
||||
def __complex__(self):
|
||||
return complex(self.tensor)
|
||||
|
||||
def is_integer(self):
|
||||
try:
|
||||
v = self.tensor.item()
|
||||
result = int(v) == v
|
||||
except Exception:
|
||||
result = False
|
||||
return result
|
||||
|
||||
def __len__(self):
|
||||
return self.tensor.shape[0]
|
||||
|
||||
def __contains__(self, x):
|
||||
return self.tensor.__contains__(x)
|
||||
|
||||
def transpose(self, *axes):
|
||||
# np.transpose(arr, axis=None) but arr.transpose(*axes)
|
||||
return _funcs.transpose(self, axes)
|
||||
|
||||
def reshape(self, *shape, order="C"):
|
||||
# arr.reshape(shape) and arr.reshape(*shape)
|
||||
return _funcs.reshape(self, shape, order=order)
|
||||
|
||||
def sort(self, axis=-1, kind=None, order=None):
|
||||
# ndarray.sort works in-place
|
||||
_funcs.copyto(self, _funcs.sort(self, axis, kind, order))
|
||||
|
||||
def item(self, *args):
|
||||
# Mimic NumPy's implementation with three special cases (no arguments,
|
||||
# a flat index and a multi-index):
|
||||
# https://github.com/numpy/numpy/blob/main/numpy/_core/src/multiarray/methods.c#L702
|
||||
if args == ():
|
||||
return self.tensor.item()
|
||||
elif len(args) == 1:
|
||||
# int argument
|
||||
return self.ravel()[args[0]]
|
||||
else:
|
||||
return self.__getitem__(args)
|
||||
|
||||
def __getitem__(self, index):
|
||||
tensor = self.tensor
|
||||
|
||||
def neg_step(i, s):
|
||||
if not (isinstance(s, slice) and s.step is not None and s.step < 0):
|
||||
return s
|
||||
|
||||
nonlocal tensor
|
||||
tensor = torch.flip(tensor, (i,))
|
||||
|
||||
# Account for the fact that a slice includes the start but not the end
|
||||
if not (isinstance(s.start, int) or s.start is None):
|
||||
raise AssertionError(
|
||||
f"slice start must be int or None, got {type(s.start).__name__}"
|
||||
)
|
||||
if not (isinstance(s.stop, int) or s.stop is None):
|
||||
raise AssertionError(
|
||||
f"slice stop must be int or None, got {type(s.stop).__name__}"
|
||||
)
|
||||
start = s.stop + 1 if s.stop else None
|
||||
stop = s.start + 1 if s.start else None
|
||||
|
||||
return slice(start, stop, -s.step)
|
||||
|
||||
if isinstance(index, Sequence):
|
||||
index = type(index)(neg_step(i, s) for i, s in enumerate(index))
|
||||
else:
|
||||
index = neg_step(0, index)
|
||||
index = _util.ndarrays_to_tensors(index)
|
||||
index = _upcast_int_indices(index)
|
||||
# Apply NumPy-compatible indexing conversion
|
||||
index = _numpy_compatible_indexing(index)
|
||||
# Apply NumPy-compatible empty ellipsis behavior
|
||||
index, maybe_squeeze, _ = _numpy_empty_ellipsis_patch(index, tensor.ndim)
|
||||
return maybe_squeeze(ndarray(tensor.__getitem__(index)))
|
||||
|
||||
def __setitem__(self, index, value):
|
||||
index = _util.ndarrays_to_tensors(index)
|
||||
index = _upcast_int_indices(index)
|
||||
# Apply NumPy-compatible indexing conversion
|
||||
index = _numpy_compatible_indexing(index)
|
||||
# Apply NumPy-compatible empty ellipsis behavior
|
||||
index, _, maybe_unsqueeze = _numpy_empty_ellipsis_patch(index, self.tensor.ndim)
|
||||
|
||||
if not _dtypes_impl.is_scalar(value):
|
||||
value = normalize_array_like(value)
|
||||
value = _util.cast_if_needed(value, self.tensor.dtype)
|
||||
|
||||
return self.tensor.__setitem__(index, maybe_unsqueeze(value))
|
||||
|
||||
take = _funcs.take
|
||||
put = _funcs.put
|
||||
|
||||
def __dlpack__(self, *, stream=None):
|
||||
return self.tensor.__dlpack__(stream=stream)
|
||||
|
||||
def __dlpack_device__(self):
|
||||
return self.tensor.__dlpack_device__()
|
||||
|
||||
|
||||
def _tolist(obj):
|
||||
"""Recursively convert tensors into lists."""
|
||||
a1 = []
|
||||
for elem in obj:
|
||||
if isinstance(elem, (list, tuple)):
|
||||
elem = _tolist(elem)
|
||||
if isinstance(elem, ndarray):
|
||||
a1.append(elem.tensor.tolist())
|
||||
else:
|
||||
a1.append(elem)
|
||||
return a1
|
||||
|
||||
|
||||
# This is the ideally the only place which talks to ndarray directly.
|
||||
# The rest goes through asarray (preferred) or array.
|
||||
|
||||
|
||||
def array(obj, dtype=None, *, copy=True, order="K", subok=False, ndmin=0, like=None):
|
||||
if subok is not False:
|
||||
raise NotImplementedError("'subok' parameter is not supported.")
|
||||
if like is not None:
|
||||
raise NotImplementedError("'like' parameter is not supported.")
|
||||
if order != "K":
|
||||
raise NotImplementedError
|
||||
|
||||
# a happy path
|
||||
if (
|
||||
isinstance(obj, ndarray)
|
||||
and copy is False
|
||||
and dtype is None
|
||||
and ndmin <= obj.ndim
|
||||
):
|
||||
return obj
|
||||
|
||||
if isinstance(obj, (list, tuple)):
|
||||
# FIXME and they have the same dtype, device, etc
|
||||
if obj and all(isinstance(x, torch.Tensor) for x in obj):
|
||||
# list of arrays: *under torch.Dynamo* these are FakeTensors
|
||||
obj = torch.stack(obj)
|
||||
else:
|
||||
# XXX: remove tolist
|
||||
# lists of ndarrays: [1, [2, 3], ndarray(4)] convert to lists of lists
|
||||
obj = _tolist(obj)
|
||||
|
||||
# is obj an ndarray already?
|
||||
if isinstance(obj, ndarray):
|
||||
obj = obj.tensor
|
||||
|
||||
# is a specific dtype requested?
|
||||
torch_dtype = None
|
||||
if dtype is not None:
|
||||
torch_dtype = _dtypes.dtype(dtype).torch_dtype
|
||||
|
||||
tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin)
|
||||
return ndarray(tensor)
|
||||
|
||||
|
||||
def asarray(a, dtype=None, order="K", *, like=None):
|
||||
return array(a, dtype=dtype, order=order, like=like, copy=False, ndmin=0)
|
||||
|
||||
|
||||
def ascontiguousarray(a, dtype=None, *, like=None):
|
||||
arr = asarray(a, dtype=dtype, like=like)
|
||||
if not arr.tensor.is_contiguous():
|
||||
arr.tensor = arr.tensor.contiguous()
|
||||
return arr
|
||||
|
||||
|
||||
def from_dlpack(x, /):
|
||||
t = torch.from_dlpack(x)
|
||||
return ndarray(t)
|
||||
|
||||
|
||||
def _extract_dtype(entry):
|
||||
try:
|
||||
dty = _dtypes.dtype(entry)
|
||||
except Exception:
|
||||
dty = asarray(entry).dtype
|
||||
return dty
|
||||
|
||||
|
||||
def can_cast(from_, to, casting="safe"):
|
||||
from_ = _extract_dtype(from_)
|
||||
to_ = _extract_dtype(to)
|
||||
|
||||
return _dtypes_impl.can_cast_impl(from_.torch_dtype, to_.torch_dtype, casting)
|
||||
|
||||
|
||||
def result_type(*arrays_and_dtypes):
|
||||
tensors = []
|
||||
for entry in arrays_and_dtypes:
|
||||
try:
|
||||
t = asarray(entry).tensor
|
||||
except (RuntimeError, ValueError, TypeError):
|
||||
dty = _dtypes.dtype(entry)
|
||||
t = torch.empty(1, dtype=dty.torch_dtype)
|
||||
tensors.append(t)
|
||||
|
||||
torch_dtype = _dtypes_impl.result_type_impl(*tensors)
|
||||
return _dtypes.dtype(torch_dtype)
|
||||
@@ -0,0 +1,267 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
""" "Normalize" arguments: convert array_likes to tensors, dtypes to torch dtypes and so on."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
import operator
|
||||
import types
|
||||
import typing
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes, _dtypes_impl, _util
|
||||
|
||||
|
||||
ArrayLike = typing.TypeVar("ArrayLike")
|
||||
Scalar = int | float | complex | bool
|
||||
ArrayLikeOrScalar = ArrayLike | Scalar
|
||||
|
||||
DTypeLike = typing.TypeVar("DTypeLike")
|
||||
AxisLike = typing.TypeVar("AxisLike")
|
||||
NDArray = typing.TypeVar("NDArray")
|
||||
CastingModes = typing.TypeVar("CastingModes")
|
||||
KeepDims = typing.TypeVar("KeepDims")
|
||||
|
||||
# OutArray is to annotate the out= array argument.
|
||||
#
|
||||
# This one is special is several respects:
|
||||
# First, It needs to be an NDArray, and we need to preserve the `result is out`
|
||||
# semantics. Therefore, we cannot just extract the Tensor from the out array.
|
||||
# So we never pass the out array to implementer functions and handle it in the
|
||||
# `normalizer` below.
|
||||
# Second, the out= argument can be either keyword or positional argument, and
|
||||
# as a positional arg, it can be anywhere in the signature.
|
||||
# To handle all this, we define a special `OutArray` annotation and dispatch on it.
|
||||
#
|
||||
OutArray = typing.TypeVar("OutArray")
|
||||
|
||||
NotImplementedType = typing.TypeVar(
|
||||
"NotImplementedType", bound=types.NotImplementedType
|
||||
)
|
||||
|
||||
|
||||
def normalize_array_like(x, parm=None): # codespell:ignore
|
||||
from ._ndarray import asarray
|
||||
|
||||
return asarray(x).tensor
|
||||
|
||||
|
||||
def normalize_array_like_or_scalar(x, parm=None): # codespell:ignore
|
||||
if _dtypes_impl.is_scalar_or_symbolic(x):
|
||||
return x
|
||||
return normalize_array_like(x, parm) # codespell:ignore
|
||||
|
||||
|
||||
def normalize_optional_array_like_or_scalar(x, parm=None): # codespell:ignore
|
||||
if x is None:
|
||||
return None
|
||||
return normalize_array_like_or_scalar(x, parm) # codespell:ignore
|
||||
|
||||
|
||||
def normalize_optional_array_like(x, parm=None): # codespell:ignore
|
||||
# This explicit normalizer is needed because otherwise normalize_array_like
|
||||
# does not run for a parameter annotated as Optional[ArrayLike]
|
||||
return None if x is None else normalize_array_like(x, parm) # codespell:ignore
|
||||
|
||||
|
||||
def normalize_seq_array_like(x, parm=None): # codespell:ignore
|
||||
return tuple(normalize_array_like(value) for value in x)
|
||||
|
||||
|
||||
def normalize_dtype(dtype, parm=None): # codespell:ignore
|
||||
# cf _decorators.dtype_to_torch
|
||||
torch_dtype = None
|
||||
if dtype is not None:
|
||||
dtype = _dtypes.dtype(dtype)
|
||||
torch_dtype = dtype.torch_dtype
|
||||
return torch_dtype
|
||||
|
||||
|
||||
def normalize_not_implemented(arg, parm): # codespell:ignore
|
||||
if arg != parm.default: # codespell:ignore
|
||||
raise NotImplementedError(
|
||||
f"'{parm.name}' parameter is not supported." # codespell:ignore
|
||||
)
|
||||
|
||||
|
||||
def normalize_axis_like(arg, parm=None): # codespell:ignore
|
||||
from ._ndarray import ndarray
|
||||
|
||||
if isinstance(arg, ndarray):
|
||||
arg = operator.index(arg)
|
||||
return arg
|
||||
|
||||
|
||||
def normalize_ndarray(arg, parm=None): # codespell:ignore
|
||||
# check the arg is an ndarray, extract its tensor attribute
|
||||
if arg is None:
|
||||
return arg
|
||||
|
||||
from ._ndarray import ndarray
|
||||
|
||||
if not isinstance(arg, ndarray):
|
||||
raise TypeError(f"'{parm.name}' must be an array") # codespell:ignore
|
||||
return arg.tensor
|
||||
|
||||
|
||||
def normalize_outarray(arg, parm=None): # codespell:ignore
|
||||
# almost normalize_ndarray, only return the array, not its tensor
|
||||
if arg is None:
|
||||
return arg
|
||||
from ._ndarray import ndarray
|
||||
|
||||
# Dynamo can pass torch tensors as out arguments,
|
||||
# wrap it in an ndarray before processing
|
||||
if isinstance(arg, torch.Tensor):
|
||||
arg = ndarray(arg)
|
||||
|
||||
if not isinstance(arg, ndarray):
|
||||
raise TypeError(f"'{parm.name}' must be an array") # codespell:ignore
|
||||
return arg
|
||||
|
||||
|
||||
def normalize_casting(arg, parm=None): # codespell:ignore
|
||||
if arg not in ["no", "equiv", "safe", "same_kind", "unsafe"]:
|
||||
raise ValueError(
|
||||
f"casting must be one of 'no', 'equiv', 'safe', 'same_kind', or 'unsafe' (got '{arg}')"
|
||||
)
|
||||
return arg
|
||||
|
||||
|
||||
normalizers = {
|
||||
"ArrayLike": normalize_array_like,
|
||||
"ArrayLikeOrScalar": normalize_array_like_or_scalar,
|
||||
"Optional[ArrayLike]": normalize_optional_array_like,
|
||||
"ArrayLike | None": normalize_optional_array_like,
|
||||
"Sequence[ArrayLike]": normalize_seq_array_like,
|
||||
"Optional[ArrayLikeOrScalar]": normalize_optional_array_like_or_scalar,
|
||||
"ArrayLikeOrScalar | None": normalize_optional_array_like_or_scalar,
|
||||
"Optional[NDArray]": normalize_ndarray,
|
||||
"NDArray | None": normalize_ndarray,
|
||||
"Optional[OutArray]": normalize_outarray,
|
||||
"OutArray | None": normalize_outarray,
|
||||
"NDArray": normalize_ndarray,
|
||||
"Optional[DTypeLike]": normalize_dtype,
|
||||
"DTypeLike | None": normalize_dtype,
|
||||
"AxisLike": normalize_axis_like,
|
||||
"NotImplementedType": normalize_not_implemented,
|
||||
"Optional[CastingModes]": normalize_casting,
|
||||
"CastingModes | None": normalize_casting,
|
||||
}
|
||||
|
||||
|
||||
def maybe_normalize(arg, parm): # codespell:ignore
|
||||
"""Normalize arg if a normalizer is registered."""
|
||||
normalizer = normalizers.get(parm.annotation) # codespell:ignore
|
||||
return normalizer(arg, parm) if normalizer else arg # codespell:ignore
|
||||
|
||||
|
||||
# ### Return value helpers ###
|
||||
|
||||
|
||||
def maybe_copy_to(out, result, promote_scalar_result=False):
|
||||
# NB: here out is either an ndarray or None
|
||||
if out is None:
|
||||
return result
|
||||
elif isinstance(result, torch.Tensor):
|
||||
if result.shape != out.shape:
|
||||
can_fit = result.numel() == 1 and out.ndim == 0
|
||||
if promote_scalar_result and can_fit:
|
||||
result = result.squeeze()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Bad size of the out array: out.shape = {out.shape}"
|
||||
f" while result.shape = {result.shape}."
|
||||
)
|
||||
out.tensor.copy_(result)
|
||||
return out
|
||||
elif isinstance(result, (tuple, list)):
|
||||
return type(result)(
|
||||
maybe_copy_to(o, r, promote_scalar_result) for o, r in zip(out, result)
|
||||
)
|
||||
else:
|
||||
raise AssertionError # We should never hit this path
|
||||
|
||||
|
||||
def wrap_tensors(result):
|
||||
from ._ndarray import ndarray
|
||||
|
||||
if isinstance(result, torch.Tensor):
|
||||
return ndarray(result)
|
||||
elif isinstance(result, (tuple, list)):
|
||||
result = type(result)(wrap_tensors(x) for x in result)
|
||||
return result
|
||||
|
||||
|
||||
def array_or_scalar(values, py_type=float, return_scalar=False):
|
||||
if return_scalar:
|
||||
return py_type(values.item())
|
||||
else:
|
||||
from ._ndarray import ndarray
|
||||
|
||||
return ndarray(values)
|
||||
|
||||
|
||||
# ### The main decorator to normalize arguments / postprocess the output ###
|
||||
|
||||
|
||||
def normalizer(_func=None, *, promote_scalar_result=False):
|
||||
def normalizer_inner(func):
|
||||
@functools.wraps(func)
|
||||
def wrapped(*args, **kwds):
|
||||
sig = inspect.signature(func)
|
||||
params = sig.parameters
|
||||
first_param = next(iter(params.values()))
|
||||
|
||||
# NumPy's API does not have positional args before variadic positional args
|
||||
if first_param.kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
args = [maybe_normalize(arg, first_param) for arg in args]
|
||||
else:
|
||||
# NB: extra unknown arguments: pass through, will raise in func(*args) below
|
||||
args = (
|
||||
tuple(
|
||||
maybe_normalize(arg, parm) # codespell:ignore
|
||||
for arg, parm in zip(args, params.values()) # codespell:ignore
|
||||
)
|
||||
+ args[len(params.values()) :]
|
||||
)
|
||||
|
||||
kwds = {
|
||||
name: maybe_normalize(arg, params[name]) if name in params else arg
|
||||
for name, arg in kwds.items()
|
||||
}
|
||||
|
||||
result = func(*args, **kwds)
|
||||
|
||||
# keepdims
|
||||
bound_args = None
|
||||
if "keepdims" in params and params["keepdims"].annotation == "KeepDims":
|
||||
# keepdims can be in any position so we need sig.bind
|
||||
bound_args = sig.bind(*args, **kwds).arguments
|
||||
if bound_args.get("keepdims", False):
|
||||
# In this case the first arg is the initial tensor and
|
||||
# the second arg is (optionally) the axis
|
||||
tensor = args[0]
|
||||
axis = bound_args.get("axis")
|
||||
result = _util.apply_keepdims(result, axis, tensor.ndim)
|
||||
|
||||
# out
|
||||
if "out" in params:
|
||||
# out can be in any position so we need sig.bind
|
||||
if bound_args is None:
|
||||
bound_args = sig.bind(*args, **kwds).arguments
|
||||
out = bound_args.get("out")
|
||||
result = maybe_copy_to(out, result, promote_scalar_result)
|
||||
result = wrap_tensors(result)
|
||||
|
||||
return result
|
||||
|
||||
return wrapped
|
||||
|
||||
if _func is None:
|
||||
return normalizer_inner
|
||||
else:
|
||||
return normalizer_inner(_func)
|
||||
@@ -0,0 +1,463 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Implementation of reduction operations, to be wrapped into arrays, dtypes etc
|
||||
in the 'public' layer.
|
||||
|
||||
Anything here only deals with torch objects, e.g. "dtype" is a torch.dtype instance etc
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl, _util
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._normalizations import (
|
||||
ArrayLike,
|
||||
AxisLike,
|
||||
DTypeLike,
|
||||
KeepDims,
|
||||
NotImplementedType,
|
||||
OutArray,
|
||||
)
|
||||
|
||||
|
||||
def _deco_axis_expand(func):
|
||||
"""
|
||||
Generically handle axis arguments in reductions.
|
||||
axis is *always* the 2nd arg in the function so no need to have a look at its signature
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapped(a, axis=None, *args, **kwds):
|
||||
if axis is not None:
|
||||
axis = _util.normalize_axis_tuple(axis, a.ndim)
|
||||
|
||||
if axis == ():
|
||||
# So we insert a length-one axis and run the reduction along it.
|
||||
# We cannot return a.clone() as this would sidestep the checks inside the function
|
||||
newshape = _util.expand_shape(a.shape, axis=0)
|
||||
a = a.reshape(newshape)
|
||||
axis = (0,)
|
||||
|
||||
return func(a, axis, *args, **kwds)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _atleast_float(dtype, other_dtype):
|
||||
"""Return a dtype that is real or complex floating-point.
|
||||
|
||||
For inputs that are boolean or integer dtypes, this returns the default
|
||||
float dtype; inputs that are complex get converted to the default complex
|
||||
dtype; real floating-point dtypes (`float*`) get passed through unchanged
|
||||
"""
|
||||
if dtype is None:
|
||||
dtype = other_dtype
|
||||
if not (dtype.is_floating_point or dtype.is_complex):
|
||||
return _dtypes_impl.default_dtypes().float_dtype
|
||||
return dtype
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims: KeepDims = False):
|
||||
return a.count_nonzero(axis)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def argmax(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
keepdims: KeepDims = False,
|
||||
):
|
||||
if a.is_complex():
|
||||
raise NotImplementedError(f"argmax with dtype={a.dtype}.")
|
||||
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
|
||||
if a.dtype == torch.bool:
|
||||
# RuntimeError: "argmax_cpu" not implemented for 'Bool'
|
||||
a = a.to(torch.uint8)
|
||||
|
||||
return torch.argmax(a, axis)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def argmin(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
keepdims: KeepDims = False,
|
||||
):
|
||||
if a.is_complex():
|
||||
raise NotImplementedError(f"argmin with dtype={a.dtype}.")
|
||||
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
|
||||
if a.dtype == torch.bool:
|
||||
# RuntimeError: "argmin_cpu" not implemented for 'Bool'
|
||||
a = a.to(torch.uint8)
|
||||
|
||||
return torch.argmin(a, axis)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def any(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
axis_kw = {} if axis is None else {"dim": axis}
|
||||
return torch.any(a, **axis_kw)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def all(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
axis_kw = {} if axis is None else {"dim": axis}
|
||||
return torch.all(a, **axis_kw)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def amax(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
initial: NotImplementedType = None,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
if a.is_complex():
|
||||
raise NotImplementedError(f"amax with dtype={a.dtype}")
|
||||
|
||||
return a.amax(axis)
|
||||
|
||||
|
||||
max = amax
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def amin(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
initial: NotImplementedType = None,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
if a.is_complex():
|
||||
raise NotImplementedError(f"amin with dtype={a.dtype}")
|
||||
|
||||
return a.amin(axis)
|
||||
|
||||
|
||||
min = amin
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def ptp(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
):
|
||||
return a.amax(axis) - a.amin(axis)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def sum(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
initial: NotImplementedType = None,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
if dtype is not None and not isinstance(dtype, torch.dtype):
|
||||
raise AssertionError(
|
||||
f"dtype must be None or a torch.dtype, got {type(dtype).__name__}"
|
||||
)
|
||||
|
||||
if dtype == torch.bool:
|
||||
dtype = _dtypes_impl.default_dtypes().int_dtype
|
||||
|
||||
axis_kw = {} if axis is None else {"dim": axis}
|
||||
return a.sum(dtype=dtype, **axis_kw)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def prod(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
initial: NotImplementedType = None,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
|
||||
if dtype == torch.bool:
|
||||
dtype = _dtypes_impl.default_dtypes().int_dtype
|
||||
|
||||
axis_kw = {} if axis is None else {"dim": axis}
|
||||
return a.prod(dtype=dtype, **axis_kw)
|
||||
|
||||
|
||||
product = prod
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def mean(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
dtype = _atleast_float(dtype, a.dtype)
|
||||
|
||||
axis_kw = {} if axis is None else {"dim": axis}
|
||||
result = a.mean(dtype=dtype, **axis_kw)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def std(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
ddof=0,
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
in_dtype = dtype
|
||||
dtype = _atleast_float(dtype, a.dtype)
|
||||
tensor = _util.cast_if_needed(a, dtype)
|
||||
result = tensor.std(dim=axis, correction=ddof)
|
||||
return _util.cast_if_needed(result, in_dtype)
|
||||
|
||||
|
||||
@_deco_axis_expand
|
||||
def var(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
ddof=0,
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
where: NotImplementedType = None,
|
||||
):
|
||||
in_dtype = dtype
|
||||
dtype = _atleast_float(dtype, a.dtype)
|
||||
tensor = _util.cast_if_needed(a, dtype)
|
||||
result = tensor.var(dim=axis, correction=ddof)
|
||||
return _util.cast_if_needed(result, in_dtype)
|
||||
|
||||
|
||||
# cumsum / cumprod are almost reductions:
|
||||
# 1. no keepdims
|
||||
# 2. axis=None flattens
|
||||
|
||||
|
||||
def cumsum(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
):
|
||||
if dtype == torch.bool:
|
||||
dtype = _dtypes_impl.default_dtypes().int_dtype
|
||||
if dtype is None:
|
||||
dtype = a.dtype
|
||||
|
||||
(a,), axis = _util.axis_none_flatten(a, axis=axis)
|
||||
axis = _util.normalize_axis_index(axis, a.ndim)
|
||||
|
||||
return a.cumsum(axis=axis, dtype=dtype)
|
||||
|
||||
|
||||
def cumprod(
|
||||
a: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
dtype: DTypeLike | None = None,
|
||||
out: OutArray | None = None,
|
||||
):
|
||||
if dtype == torch.bool:
|
||||
dtype = _dtypes_impl.default_dtypes().int_dtype
|
||||
if dtype is None:
|
||||
dtype = a.dtype
|
||||
|
||||
(a,), axis = _util.axis_none_flatten(a, axis=axis)
|
||||
axis = _util.normalize_axis_index(axis, a.ndim)
|
||||
|
||||
return a.cumprod(axis=axis, dtype=dtype)
|
||||
|
||||
|
||||
cumproduct = cumprod
|
||||
|
||||
|
||||
def average(
|
||||
a: ArrayLike,
|
||||
axis=None,
|
||||
weights: ArrayLike = None,
|
||||
returned=False,
|
||||
*,
|
||||
keepdims=False,
|
||||
):
|
||||
if weights is None:
|
||||
result = mean(a, axis=axis)
|
||||
wsum = torch.as_tensor(a.numel() / result.numel(), dtype=result.dtype)
|
||||
else:
|
||||
if not a.dtype.is_floating_point:
|
||||
a = a.double()
|
||||
|
||||
# axis & weights
|
||||
if a.shape != weights.shape:
|
||||
if axis is None:
|
||||
raise TypeError(
|
||||
"Axis must be specified when shapes of a and weights differ."
|
||||
)
|
||||
if weights.ndim != 1:
|
||||
raise TypeError(
|
||||
"1D weights expected when shapes of a and weights differ."
|
||||
)
|
||||
if weights.shape[0] != a.shape[axis]:
|
||||
raise ValueError(
|
||||
"Length of weights not compatible with specified axis."
|
||||
)
|
||||
|
||||
# setup weight to broadcast along axis
|
||||
weights = torch.broadcast_to(weights, (a.ndim - 1) * (1,) + weights.shape)
|
||||
weights = weights.swapaxes(-1, axis)
|
||||
|
||||
# do the work
|
||||
result_dtype = _dtypes_impl.result_type_impl(a, weights)
|
||||
numerator = sum(a * weights, axis, dtype=result_dtype)
|
||||
wsum = sum(weights, axis, dtype=result_dtype)
|
||||
result = numerator / wsum
|
||||
|
||||
# We process keepdims manually because the decorator does not deal with variadic returns
|
||||
if keepdims:
|
||||
result = _util.apply_keepdims(result, axis, a.ndim)
|
||||
|
||||
if returned:
|
||||
if wsum.shape != result.shape:
|
||||
wsum = torch.broadcast_to(wsum, result.shape).clone()
|
||||
return result, wsum
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
# Not using deco_axis_expand as it assumes that axis is the second arg
|
||||
def quantile(
|
||||
a: ArrayLike,
|
||||
q: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
overwrite_input=False,
|
||||
method="linear",
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
interpolation: NotImplementedType = None,
|
||||
):
|
||||
if overwrite_input:
|
||||
# raise NotImplementedError("overwrite_input in quantile not implemented.")
|
||||
# NumPy documents that `overwrite_input` MAY modify inputs:
|
||||
# https://numpy.org/doc/stable/reference/generated/numpy.percentile.html#numpy-percentile
|
||||
# Here we choose to work out-of-place because why not.
|
||||
pass
|
||||
|
||||
if not a.dtype.is_floating_point:
|
||||
dtype = _dtypes_impl.default_dtypes().float_dtype
|
||||
a = a.to(dtype)
|
||||
|
||||
# edge case: torch.quantile only supports float32 and float64
|
||||
if a.dtype == torch.float16:
|
||||
a = a.to(torch.float32)
|
||||
|
||||
if axis is None:
|
||||
a = a.flatten()
|
||||
q = q.flatten()
|
||||
axis = (0,)
|
||||
else:
|
||||
axis = _util.normalize_axis_tuple(axis, a.ndim)
|
||||
|
||||
# FIXME(Mario) Doesn't np.quantile accept a tuple?
|
||||
# torch.quantile does accept a number. If we don't want to implement the tuple behaviour
|
||||
# (it's deffo low prio) change `normalize_axis_tuple` into a normalize_axis index above.
|
||||
axis = _util.allow_only_single_axis(axis)
|
||||
|
||||
q = _util.cast_if_needed(q, a.dtype)
|
||||
|
||||
return torch.quantile(a, q, axis=axis, interpolation=method)
|
||||
|
||||
|
||||
def percentile(
|
||||
a: ArrayLike,
|
||||
q: ArrayLike,
|
||||
axis: AxisLike = None,
|
||||
out: OutArray | None = None,
|
||||
overwrite_input=False,
|
||||
method="linear",
|
||||
keepdims: KeepDims = False,
|
||||
*,
|
||||
interpolation: NotImplementedType = None,
|
||||
):
|
||||
# np.percentile(float_tensor, 30) : q.dtype is int64 => q / 100.0 is float32
|
||||
if _dtypes_impl.python_type_for_torch(q.dtype) is int:
|
||||
q = q.to(_dtypes_impl.default_dtypes().float_dtype)
|
||||
qq = q / 100.0
|
||||
|
||||
return quantile(
|
||||
a,
|
||||
qq,
|
||||
axis=axis,
|
||||
overwrite_input=overwrite_input,
|
||||
method=method,
|
||||
keepdims=keepdims,
|
||||
interpolation=interpolation,
|
||||
)
|
||||
|
||||
|
||||
def median(
|
||||
a: ArrayLike,
|
||||
axis=None,
|
||||
out: OutArray | None = None,
|
||||
overwrite_input=False,
|
||||
keepdims: KeepDims = False,
|
||||
):
|
||||
return quantile(
|
||||
a,
|
||||
torch.as_tensor(0.5),
|
||||
axis=axis,
|
||||
overwrite_input=overwrite_input,
|
||||
out=out,
|
||||
keepdims=keepdims,
|
||||
)
|
||||
@@ -0,0 +1,332 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from . import _binary_ufuncs_impl, _dtypes_impl, _unary_ufuncs_impl, _util
|
||||
from ._normalizations import (
|
||||
ArrayLike,
|
||||
ArrayLikeOrScalar,
|
||||
CastingModes,
|
||||
DTypeLike,
|
||||
normalizer,
|
||||
NotImplementedType,
|
||||
OutArray,
|
||||
)
|
||||
|
||||
|
||||
def _ufunc_postprocess(result, out, casting):
|
||||
if out is not None:
|
||||
result = _util.typecast_tensor(result, out.dtype.torch_dtype, casting)
|
||||
result = torch.broadcast_to(result, out.shape)
|
||||
return result
|
||||
|
||||
|
||||
# ############# Binary ufuncs ######################
|
||||
|
||||
_binary = [
|
||||
name
|
||||
for name in dir(_binary_ufuncs_impl)
|
||||
if not name.startswith("_") and name not in ["torch", "matmul", "divmod", "ldexp"]
|
||||
]
|
||||
|
||||
|
||||
NEP50_FUNCS = (
|
||||
"add",
|
||||
"subtract",
|
||||
"multiply",
|
||||
"floor_divide",
|
||||
"true_divide",
|
||||
"divide",
|
||||
"remainder",
|
||||
"bitwise_and",
|
||||
"bitwise_or",
|
||||
"bitwise_xor",
|
||||
"bitwise_left_shift",
|
||||
"bitwise_right_shift",
|
||||
"hypot",
|
||||
"arctan2",
|
||||
"logaddexp",
|
||||
"logaddexp2",
|
||||
"heaviside",
|
||||
"copysign",
|
||||
"fmax",
|
||||
"minimum",
|
||||
"fmin",
|
||||
"maximum",
|
||||
"fmod",
|
||||
"gcd",
|
||||
"lcm",
|
||||
"pow",
|
||||
)
|
||||
|
||||
|
||||
def deco_binary_ufunc(torch_func):
|
||||
"""Common infra for binary ufuncs.
|
||||
|
||||
Normalize arguments, sort out type casting, broadcasting and delegate to
|
||||
the pytorch functions for the actual work.
|
||||
"""
|
||||
|
||||
@normalizer
|
||||
def wrapped(
|
||||
x1: ArrayLikeOrScalar,
|
||||
x2: ArrayLikeOrScalar,
|
||||
/,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
where: NotImplementedType = True,
|
||||
casting: CastingModes | None = "same_kind",
|
||||
order: NotImplementedType = "K",
|
||||
dtype: DTypeLike | None = None,
|
||||
subok: NotImplementedType = False,
|
||||
signature: NotImplementedType = None,
|
||||
extobj: NotImplementedType = None,
|
||||
):
|
||||
if dtype is not None:
|
||||
|
||||
def cast(x, dtype):
|
||||
if isinstance(x, torch.Tensor):
|
||||
return _util.typecast_tensor(x, dtype, casting)
|
||||
else:
|
||||
return torch.as_tensor(x, dtype=dtype)
|
||||
|
||||
x1 = cast(x1, dtype)
|
||||
x2 = cast(x2, dtype)
|
||||
elif isinstance(x1, torch.Tensor) and isinstance(x2, torch.Tensor):
|
||||
dtype = _dtypes_impl.result_type_impl(x1, x2)
|
||||
x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
|
||||
else:
|
||||
x1, x2 = _dtypes_impl.nep50_to_tensors(
|
||||
x1, x2, torch_func.__name__ in NEP50_FUNCS, torch_func.__name__
|
||||
)
|
||||
|
||||
result = torch_func(x1, x2)
|
||||
|
||||
return _ufunc_postprocess(result, out, casting)
|
||||
|
||||
wrapped.__qualname__ = torch_func.__name__
|
||||
wrapped.__name__ = torch_func.__name__
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
# matmul's signature is _slightly_ different from other ufuncs:
|
||||
# - no where=...
|
||||
# - additional axis=..., axes=...
|
||||
# - no NEP50 scalars in or out
|
||||
@normalizer
|
||||
def matmul(
|
||||
x1: ArrayLike,
|
||||
x2: ArrayLike,
|
||||
/,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
casting: CastingModes | None = "same_kind",
|
||||
order: NotImplementedType = "K",
|
||||
dtype: DTypeLike | None = None,
|
||||
subok: NotImplementedType = False,
|
||||
signature: NotImplementedType = None,
|
||||
extobj: NotImplementedType = None,
|
||||
axes: NotImplementedType = None,
|
||||
axis: NotImplementedType = None,
|
||||
):
|
||||
if dtype is None:
|
||||
dtype = _dtypes_impl.result_type_impl(x1, x2)
|
||||
x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
|
||||
|
||||
result = _binary_ufuncs_impl.matmul(x1, x2)
|
||||
|
||||
result = _ufunc_postprocess(result, out, casting)
|
||||
return result
|
||||
|
||||
|
||||
# ldexp casting is special : the dtype of the result == dtype of the 1st arg
|
||||
@normalizer
|
||||
def ldexp(
|
||||
x1: ArrayLikeOrScalar,
|
||||
x2: ArrayLikeOrScalar,
|
||||
/,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
where: NotImplementedType = True,
|
||||
casting: CastingModes | None = "same_kind",
|
||||
order: NotImplementedType = "K",
|
||||
dtype: DTypeLike | None = None,
|
||||
subok: NotImplementedType = False,
|
||||
signature: NotImplementedType = None,
|
||||
extobj: NotImplementedType = None,
|
||||
):
|
||||
if dtype is not None:
|
||||
if isinstance(x1, torch.Tensor):
|
||||
x1 = _util.typecast_tensor(x1, dtype, casting)
|
||||
else:
|
||||
x1 = torch.as_tensor(x1, dtype=dtype)
|
||||
else:
|
||||
if not isinstance(x1, torch.Tensor):
|
||||
x1 = torch.as_tensor(x1)
|
||||
x1 = _util.cast_int_to_float(x1)
|
||||
|
||||
x2 = torch.as_tensor(x2)
|
||||
# the second arg must be integer
|
||||
if _dtypes_impl._category(x2.dtype) != 1:
|
||||
raise ValueError("ldexp 2nd arg must be integer")
|
||||
|
||||
result = _binary_ufuncs_impl.ldexp(x1, x2)
|
||||
|
||||
if x1.dtype == torch.float16:
|
||||
# torch.ldexp(f16, int) -> f32, undo it
|
||||
result = result.to(torch.float16)
|
||||
|
||||
return _ufunc_postprocess(result, out, casting)
|
||||
|
||||
|
||||
# nin=2, nout=2
|
||||
@normalizer
|
||||
def divmod(
|
||||
x1: ArrayLike,
|
||||
x2: ArrayLike,
|
||||
out1: OutArray | None = None,
|
||||
out2: OutArray | None = None,
|
||||
/,
|
||||
out: tuple[OutArray | None, OutArray | None] = (None, None),
|
||||
*,
|
||||
where: NotImplementedType = True,
|
||||
casting: CastingModes | None = "same_kind",
|
||||
order: NotImplementedType = "K",
|
||||
dtype: DTypeLike | None = None,
|
||||
subok: NotImplementedType = False,
|
||||
signature: NotImplementedType = None,
|
||||
extobj: NotImplementedType = None,
|
||||
):
|
||||
# make sure we either have no out arrays at all, or there is either
|
||||
# out1, out2, or out=tuple, but not both
|
||||
num_outs = sum(x is not None for x in [out1, out2])
|
||||
if num_outs == 1:
|
||||
raise ValueError("both out1 and out2 need to be provided")
|
||||
elif num_outs == 2:
|
||||
o1, o2 = out
|
||||
if o1 is not None or o2 is not None:
|
||||
raise TypeError(
|
||||
"cannot specify 'out' as both a positional and keyword argument"
|
||||
)
|
||||
else:
|
||||
out1, out2 = out
|
||||
|
||||
if dtype is None:
|
||||
dtype = _dtypes_impl.result_type_impl(x1, x2)
|
||||
x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
|
||||
|
||||
quot, rem = _binary_ufuncs_impl.divmod(x1, x2)
|
||||
|
||||
quot = _ufunc_postprocess(quot, out1, casting)
|
||||
rem = _ufunc_postprocess(rem, out2, casting)
|
||||
return quot, rem
|
||||
|
||||
|
||||
#
|
||||
# Attach ufuncs to this module, for a further export to the public namespace in __init__.py
|
||||
#
|
||||
for name in _binary:
|
||||
ufunc = getattr(_binary_ufuncs_impl, name)
|
||||
vars()[name] = deco_binary_ufunc(ufunc)
|
||||
|
||||
|
||||
def modf(x, /, *args, **kwds):
|
||||
quot, rem = divmod(x, 1, *args, **kwds)
|
||||
return rem, quot
|
||||
|
||||
|
||||
_binary = _binary + ["divmod", "modf", "matmul", "ldexp"]
|
||||
|
||||
|
||||
# ############# Unary ufuncs ######################
|
||||
|
||||
|
||||
_unary = [
|
||||
name
|
||||
for name in dir(_unary_ufuncs_impl)
|
||||
if not name.startswith("_") and name != "torch"
|
||||
]
|
||||
|
||||
|
||||
# these are ufunc(int) -> float
|
||||
_fp_unary = [
|
||||
"arccos",
|
||||
"arccosh",
|
||||
"arcsin",
|
||||
"arcsinh",
|
||||
"arctan",
|
||||
"arctanh",
|
||||
"cbrt",
|
||||
"cos",
|
||||
"cosh",
|
||||
"deg2rad",
|
||||
"degrees",
|
||||
"exp",
|
||||
"exp2",
|
||||
"expm1",
|
||||
"log",
|
||||
"log10",
|
||||
"log1p",
|
||||
"log2",
|
||||
"rad2deg",
|
||||
"radians",
|
||||
"reciprocal",
|
||||
"sin",
|
||||
"sinh",
|
||||
"sqrt",
|
||||
"square",
|
||||
"tan",
|
||||
"tanh",
|
||||
"trunc",
|
||||
]
|
||||
|
||||
|
||||
def deco_unary_ufunc(torch_func):
|
||||
"""Common infra for unary ufuncs.
|
||||
|
||||
Normalize arguments, sort out type casting, broadcasting and delegate to
|
||||
the pytorch functions for the actual work.
|
||||
"""
|
||||
|
||||
@normalizer
|
||||
def wrapped(
|
||||
x: ArrayLike,
|
||||
/,
|
||||
out: OutArray | None = None,
|
||||
*,
|
||||
where=True,
|
||||
casting: CastingModes | None = "same_kind",
|
||||
order="K",
|
||||
dtype: DTypeLike | None = None,
|
||||
subok: NotImplementedType = False,
|
||||
signature=None,
|
||||
extobj=None,
|
||||
):
|
||||
if dtype is not None:
|
||||
x = _util.typecast_tensor(x, dtype, casting)
|
||||
|
||||
if torch_func.__name__ in _fp_unary:
|
||||
x = _util.cast_int_to_float(x)
|
||||
|
||||
result = torch_func(x)
|
||||
result = _ufunc_postprocess(result, out, casting)
|
||||
return result
|
||||
|
||||
wrapped.__qualname__ = torch_func.__name__
|
||||
wrapped.__name__ = torch_func.__name__
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
#
|
||||
# Attach ufuncs to this module, for a further export to the public namespace in __init__.py
|
||||
#
|
||||
for name in _unary:
|
||||
ufunc = getattr(_unary_ufuncs_impl, name)
|
||||
vars()[name] = deco_unary_ufunc(ufunc)
|
||||
|
||||
|
||||
__all__ = _binary + _unary # noqa: PLE0605
|
||||
@@ -0,0 +1,72 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Export torch work functions for unary ufuncs, rename/tweak to match numpy.
|
||||
This listing is further exported to public symbols in the `_numpy/_ufuncs.py` module.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import ( # noqa: F401
|
||||
absolute as fabs,
|
||||
arccos,
|
||||
arccosh,
|
||||
arcsin,
|
||||
arcsinh,
|
||||
arctan,
|
||||
arctanh,
|
||||
bitwise_not,
|
||||
bitwise_not as invert,
|
||||
ceil,
|
||||
conj_physical as conjugate,
|
||||
cos,
|
||||
cosh,
|
||||
deg2rad,
|
||||
deg2rad as radians,
|
||||
exp,
|
||||
exp2,
|
||||
expm1,
|
||||
floor,
|
||||
isfinite,
|
||||
isinf,
|
||||
isnan,
|
||||
log,
|
||||
log10,
|
||||
log1p,
|
||||
log2,
|
||||
logical_not,
|
||||
negative,
|
||||
rad2deg,
|
||||
rad2deg as degrees,
|
||||
reciprocal,
|
||||
round as fix,
|
||||
round as rint,
|
||||
sign,
|
||||
signbit,
|
||||
sin,
|
||||
sinh,
|
||||
sqrt,
|
||||
square,
|
||||
tan,
|
||||
tanh,
|
||||
trunc,
|
||||
)
|
||||
|
||||
|
||||
# special cases: torch does not export these names
|
||||
def cbrt(x):
|
||||
return torch.pow(x, 1 / 3)
|
||||
|
||||
|
||||
def positive(x):
|
||||
return +x
|
||||
|
||||
|
||||
def absolute(x):
|
||||
# work around torch.absolute not impl for bools
|
||||
if x.dtype == torch.bool:
|
||||
return x
|
||||
return torch.absolute(x)
|
||||
|
||||
|
||||
# TODO set __name__ and __qualname__
|
||||
abs = absolute
|
||||
conj = conjugate
|
||||
@@ -0,0 +1,269 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Assorted utilities, which do not need anything other then torch and stdlib."""
|
||||
|
||||
import operator
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl
|
||||
|
||||
|
||||
# https://github.com/numpy/numpy/blob/v1.23.0/numpy/distutils/misc_util.py#L497-L504
|
||||
def is_sequence(seq):
|
||||
if isinstance(seq, str):
|
||||
return False
|
||||
try:
|
||||
len(seq)
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class AxisError(ValueError, IndexError):
|
||||
pass
|
||||
|
||||
|
||||
class UFuncTypeError(TypeError, RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def cast_if_needed(tensor, dtype):
|
||||
# NB: no casting if dtype=None
|
||||
if dtype is not None and tensor.dtype != dtype:
|
||||
tensor = tensor.to(dtype)
|
||||
return tensor
|
||||
|
||||
|
||||
def cast_int_to_float(x):
|
||||
# cast integers and bools to the default float dtype
|
||||
if _dtypes_impl._category(x.dtype) < 2:
|
||||
x = x.to(_dtypes_impl.default_dtypes().float_dtype)
|
||||
return x
|
||||
|
||||
|
||||
# a replica of the version in ./numpy/numpy/core/src/multiarray/common.h
|
||||
def normalize_axis_index(ax, ndim, argname=None):
|
||||
if not (-ndim <= ax < ndim):
|
||||
raise AxisError(f"axis {ax} is out of bounds for array of dimension {ndim}")
|
||||
if ax < 0:
|
||||
ax += ndim
|
||||
return ax
|
||||
|
||||
|
||||
# from https://github.com/numpy/numpy/blob/main/numpy/core/numeric.py#L1378
|
||||
def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
|
||||
"""
|
||||
Normalizes an axis argument into a tuple of non-negative integer axes.
|
||||
|
||||
This handles shorthands such as ``1`` and converts them to ``(1,)``,
|
||||
as well as performing the handling of negative indices covered by
|
||||
`normalize_axis_index`.
|
||||
|
||||
By default, this forbids axes from being specified multiple times.
|
||||
Used internally by multi-axis-checking logic.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
axis : int, iterable of int
|
||||
The un-normalized index or indices of the axis.
|
||||
ndim : int
|
||||
The number of dimensions of the array that `axis` should be normalized
|
||||
against.
|
||||
argname : str, optional
|
||||
A prefix to put before the error message, typically the name of the
|
||||
argument.
|
||||
allow_duplicate : bool, optional
|
||||
If False, the default, disallow an axis from being specified twice.
|
||||
|
||||
Returns
|
||||
-------
|
||||
normalized_axes : tuple of int
|
||||
The normalized axis index, such that `0 <= normalized_axis < ndim`
|
||||
"""
|
||||
# Optimization to speed-up the most common cases.
|
||||
if type(axis) not in (tuple, list):
|
||||
try:
|
||||
axis = [operator.index(axis)]
|
||||
except TypeError:
|
||||
pass
|
||||
# Going via an iterator directly is slower than via list comprehension.
|
||||
axis = tuple(normalize_axis_index(ax, ndim, argname) for ax in axis)
|
||||
if not allow_duplicate and len(set(map(int, axis))) != len(axis):
|
||||
if argname:
|
||||
raise ValueError(f"repeated axis in `{argname}` argument")
|
||||
else:
|
||||
raise ValueError("repeated axis")
|
||||
return axis
|
||||
|
||||
|
||||
def allow_only_single_axis(axis):
|
||||
if axis is None:
|
||||
return axis
|
||||
if len(axis) != 1:
|
||||
raise NotImplementedError("does not handle tuple axis")
|
||||
return axis[0]
|
||||
|
||||
|
||||
def expand_shape(arr_shape, axis):
|
||||
# taken from numpy 1.23.x, expand_dims function
|
||||
if type(axis) not in (list, tuple):
|
||||
axis = (axis,)
|
||||
out_ndim = len(axis) + len(arr_shape)
|
||||
axis = normalize_axis_tuple(axis, out_ndim)
|
||||
shape_it = iter(arr_shape)
|
||||
shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)]
|
||||
return shape
|
||||
|
||||
|
||||
def apply_keepdims(tensor, axis, ndim):
|
||||
if axis is None:
|
||||
# tensor was a scalar
|
||||
shape = (1,) * ndim
|
||||
tensor = tensor.expand(shape).contiguous()
|
||||
else:
|
||||
shape = expand_shape(tensor.shape, axis)
|
||||
tensor = tensor.reshape(shape)
|
||||
return tensor
|
||||
|
||||
|
||||
def axis_none_flatten(*tensors, axis=None):
|
||||
"""Flatten the arrays if axis is None."""
|
||||
if axis is None:
|
||||
tensors = tuple(ar.flatten() for ar in tensors)
|
||||
return tensors, 0
|
||||
else:
|
||||
return tensors, axis
|
||||
|
||||
|
||||
def typecast_tensor(t, target_dtype, casting):
|
||||
"""Dtype-cast tensor to target_dtype.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
t : torch.Tensor
|
||||
The tensor to cast
|
||||
target_dtype : torch dtype object
|
||||
The array dtype to cast all tensors to
|
||||
casting : str
|
||||
The casting mode, see `np.can_cast`
|
||||
|
||||
Returns
|
||||
-------
|
||||
`torch.Tensor` of the `target_dtype` dtype
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
if the argument cannot be cast according to the `casting` rule
|
||||
|
||||
"""
|
||||
can_cast = _dtypes_impl.can_cast_impl
|
||||
|
||||
if not can_cast(t.dtype, target_dtype, casting=casting):
|
||||
raise TypeError(
|
||||
f"Cannot cast array data from {t.dtype} to"
|
||||
f" {target_dtype} according to the rule '{casting}'"
|
||||
)
|
||||
return cast_if_needed(t, target_dtype)
|
||||
|
||||
|
||||
def typecast_tensors(tensors, target_dtype, casting):
|
||||
return tuple(typecast_tensor(t, target_dtype, casting) for t in tensors)
|
||||
|
||||
|
||||
def _try_convert_to_tensor(obj):
|
||||
try:
|
||||
tensor = torch.as_tensor(obj)
|
||||
except Exception as e:
|
||||
mesg = f"failed to convert {obj} to ndarray. \nInternal error is: {str(e)}."
|
||||
raise NotImplementedError(mesg) # noqa: B904
|
||||
return tensor
|
||||
|
||||
|
||||
def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
|
||||
"""The core logic of the array(...) function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
obj : tensor_like
|
||||
The thing to coerce
|
||||
dtype : torch.dtype object or None
|
||||
Coerce to this torch dtype
|
||||
copy : bool
|
||||
Copy or not
|
||||
ndmin : int
|
||||
The results as least this many dimensions
|
||||
is_weak : bool
|
||||
Whether obj is a weakly typed python scalar.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor : torch.Tensor
|
||||
a tensor object with requested dtype, ndim and copy semantics.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This is almost a "tensor_like" coercive function. Does not handle wrapper
|
||||
ndarrays (those should be handled in the ndarray-aware layer prior to
|
||||
invoking this function).
|
||||
"""
|
||||
if isinstance(obj, torch.Tensor):
|
||||
tensor = obj
|
||||
else:
|
||||
# tensor.dtype is the pytorch default, typically float32. If obj's elements
|
||||
# are not exactly representable in float32, we've lost precision:
|
||||
# >>> torch.as_tensor(1e12).item() - 1e12
|
||||
# -4096.0
|
||||
default_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(_dtypes_impl.get_default_dtype_for(torch.float32))
|
||||
try:
|
||||
tensor = _try_convert_to_tensor(obj)
|
||||
finally:
|
||||
torch.set_default_dtype(default_dtype)
|
||||
|
||||
# type cast if requested
|
||||
tensor = cast_if_needed(tensor, dtype)
|
||||
|
||||
# adjust ndim if needed
|
||||
ndim_extra = ndmin - tensor.ndim
|
||||
if ndim_extra > 0:
|
||||
tensor = tensor.view((1,) * ndim_extra + tensor.shape)
|
||||
|
||||
# special handling for np._CopyMode
|
||||
try:
|
||||
copy = bool(copy)
|
||||
except ValueError:
|
||||
# TODO handle _CopyMode.IF_NEEDED correctly
|
||||
copy = False
|
||||
# copy if requested
|
||||
if copy:
|
||||
tensor = tensor.clone()
|
||||
|
||||
return tensor
|
||||
|
||||
|
||||
def ndarrays_to_tensors(*inputs):
|
||||
"""Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
|
||||
from ._ndarray import ndarray
|
||||
|
||||
if len(inputs) == 0:
|
||||
return ValueError()
|
||||
elif len(inputs) == 1:
|
||||
input_ = inputs[0]
|
||||
if isinstance(input_, ndarray):
|
||||
return input_.tensor
|
||||
elif isinstance(input_, tuple):
|
||||
result = []
|
||||
for sub_input in input_:
|
||||
sub_result = ndarrays_to_tensors(sub_input)
|
||||
result.append(sub_result)
|
||||
return tuple(result)
|
||||
else:
|
||||
return input_
|
||||
else:
|
||||
if not isinstance(inputs, tuple):
|
||||
raise AssertionError(
|
||||
f"Expected inputs to be a tuple, got {type(inputs).__name__}"
|
||||
)
|
||||
return ndarrays_to_tensors(inputs)
|
||||
@@ -0,0 +1,130 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl, _util
|
||||
from ._normalizations import ArrayLike, normalizer
|
||||
|
||||
|
||||
def upcast(func):
|
||||
"""NumPy fft casts inputs to 64 bit and *returns 64-bit results*."""
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapped(tensor, *args, **kwds):
|
||||
target_dtype = (
|
||||
_dtypes_impl.default_dtypes().complex_dtype
|
||||
if tensor.is_complex()
|
||||
else _dtypes_impl.default_dtypes().float_dtype
|
||||
)
|
||||
tensor = _util.cast_if_needed(tensor, target_dtype)
|
||||
return func(tensor, *args, **kwds)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def fft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.fft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def ifft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.ifft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def rfft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.rfft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def irfft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.irfft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def fftn(a: ArrayLike, s=None, axes=None, norm=None):
|
||||
return torch.fft.fftn(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def ifftn(a: ArrayLike, s=None, axes=None, norm=None):
|
||||
return torch.fft.ifftn(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def rfftn(a: ArrayLike, s=None, axes=None, norm=None):
|
||||
return torch.fft.rfftn(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def irfftn(a: ArrayLike, s=None, axes=None, norm=None):
|
||||
return torch.fft.irfftn(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def fft2(a: ArrayLike, s=None, axes=(-2, -1), norm=None):
|
||||
return torch.fft.fft2(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def ifft2(a: ArrayLike, s=None, axes=(-2, -1), norm=None):
|
||||
return torch.fft.ifft2(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def rfft2(a: ArrayLike, s=None, axes=(-2, -1), norm=None):
|
||||
return torch.fft.rfft2(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def irfft2(a: ArrayLike, s=None, axes=(-2, -1), norm=None):
|
||||
return torch.fft.irfft2(a, s, dim=axes, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def hfft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.hfft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
@upcast
|
||||
def ihfft(a: ArrayLike, n=None, axis=-1, norm=None):
|
||||
return torch.fft.ihfft(a, n, dim=axis, norm=norm)
|
||||
|
||||
|
||||
@normalizer
|
||||
def fftfreq(n, d=1.0):
|
||||
return torch.fft.fftfreq(n, d)
|
||||
|
||||
|
||||
@normalizer
|
||||
def rfftfreq(n, d=1.0):
|
||||
return torch.fft.rfftfreq(n, d)
|
||||
|
||||
|
||||
@normalizer
|
||||
def fftshift(x: ArrayLike, axes=None):
|
||||
return torch.fft.fftshift(x, axes)
|
||||
|
||||
|
||||
@normalizer
|
||||
def ifftshift(x: ArrayLike, axes=None):
|
||||
return torch.fft.ifftshift(x, axes)
|
||||
@@ -0,0 +1,243 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import math
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl, _util
|
||||
from ._normalizations import ArrayLike, KeepDims, normalizer
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
class LinAlgError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _atleast_float_1(a):
|
||||
if not (a.dtype.is_floating_point or a.dtype.is_complex):
|
||||
a = a.to(_dtypes_impl.default_dtypes().float_dtype)
|
||||
return a
|
||||
|
||||
|
||||
def _atleast_float_2(a, b):
|
||||
dtyp = _dtypes_impl.result_type_impl(a, b)
|
||||
if not (dtyp.is_floating_point or dtyp.is_complex):
|
||||
dtyp = _dtypes_impl.default_dtypes().float_dtype
|
||||
|
||||
a = _util.cast_if_needed(a, dtyp)
|
||||
b = _util.cast_if_needed(b, dtyp)
|
||||
return a, b
|
||||
|
||||
|
||||
def linalg_errors(func):
|
||||
@functools.wraps(func)
|
||||
def wrapped(*args, **kwds):
|
||||
try:
|
||||
return func(*args, **kwds)
|
||||
except torch._C._LinAlgError as e:
|
||||
raise LinAlgError(*e.args) # noqa: B904
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
# ### Matrix and vector products ###
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def matrix_power(a: ArrayLike, n):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.matrix_power(a, n)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def multi_dot(inputs: Sequence[ArrayLike], *, out=None):
|
||||
return torch.linalg.multi_dot(inputs)
|
||||
|
||||
|
||||
# ### Solving equations and inverting matrices ###
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def solve(a: ArrayLike, b: ArrayLike):
|
||||
a, b = _atleast_float_2(a, b)
|
||||
return torch.linalg.solve(a, b)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def lstsq(a: ArrayLike, b: ArrayLike, rcond=None):
|
||||
a, b = _atleast_float_2(a, b)
|
||||
# NumPy is using gelsd: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/umath_linalg.cpp#L3991
|
||||
# on CUDA, only `gels` is available though, so use it instead
|
||||
driver = "gels" if a.is_cuda or b.is_cuda else "gelsd"
|
||||
return torch.linalg.lstsq(a, b, rcond=rcond, driver=driver)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def inv(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
result = torch.linalg.inv(a)
|
||||
return result
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def pinv(a: ArrayLike, rcond=1e-15, hermitian=False):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.pinv(a, rtol=rcond, hermitian=hermitian)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def tensorsolve(a: ArrayLike, b: ArrayLike, axes=None):
|
||||
a, b = _atleast_float_2(a, b)
|
||||
return torch.linalg.tensorsolve(a, b, dims=axes)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def tensorinv(a: ArrayLike, ind=2):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.tensorinv(a, ind=ind)
|
||||
|
||||
|
||||
# ### Norms and other numbers ###
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def det(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.det(a)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def slogdet(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.slogdet(a)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def cond(x: ArrayLike, p=None):
|
||||
x = _atleast_float_1(x)
|
||||
|
||||
# check if empty
|
||||
# cf: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1744
|
||||
if x.numel() == 0 and math.prod(x.shape[-2:]) == 0:
|
||||
raise LinAlgError("cond is not defined on empty arrays")
|
||||
|
||||
result = torch.linalg.cond(x, p=p)
|
||||
|
||||
# Convert nans to infs (numpy does it in a data-dependent way, depending on
|
||||
# whether the input array has nans or not)
|
||||
# XXX: NumPy does this: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1744
|
||||
return torch.where(torch.isnan(result), float("inf"), result)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def matrix_rank(a: ArrayLike, tol=None, hermitian=False):
|
||||
a = _atleast_float_1(a)
|
||||
|
||||
if a.ndim < 2:
|
||||
return int((a != 0).any())
|
||||
|
||||
if tol is None:
|
||||
# follow https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1885
|
||||
atol = 0
|
||||
rtol = max(a.shape[-2:]) * torch.finfo(a.dtype).eps
|
||||
else:
|
||||
atol, rtol = tol, 0
|
||||
return torch.linalg.matrix_rank(a, atol=atol, rtol=rtol, hermitian=hermitian)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def norm(x: ArrayLike, ord=None, axis=None, keepdims: KeepDims = False):
|
||||
x = _atleast_float_1(x)
|
||||
return torch.linalg.norm(x, ord=ord, dim=axis)
|
||||
|
||||
|
||||
# ### Decompositions ###
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def cholesky(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.cholesky(a)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def qr(a: ArrayLike, mode="reduced"):
|
||||
a = _atleast_float_1(a)
|
||||
result = torch.linalg.qr(a, mode=mode)
|
||||
if mode == "r":
|
||||
# match NumPy
|
||||
result = result.R
|
||||
return result
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def svd(a: ArrayLike, full_matrices=True, compute_uv=True, hermitian=False):
|
||||
a = _atleast_float_1(a)
|
||||
if not compute_uv:
|
||||
return torch.linalg.svdvals(a)
|
||||
|
||||
# NB: ignore the hermitian= argument (no pytorch equivalent)
|
||||
result = torch.linalg.svd(a, full_matrices=full_matrices)
|
||||
return result
|
||||
|
||||
|
||||
# ### Eigenvalues and eigenvectors ###
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def eig(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
w, vt = torch.linalg.eig(a)
|
||||
|
||||
if not a.is_complex() and w.is_complex() and (w.imag == 0).all():
|
||||
w = w.real
|
||||
vt = vt.real
|
||||
return w, vt
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def eigh(a: ArrayLike, UPLO="L"):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.eigh(a, UPLO=UPLO)
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def eigvals(a: ArrayLike):
|
||||
a = _atleast_float_1(a)
|
||||
result = torch.linalg.eigvals(a)
|
||||
if not a.is_complex() and result.is_complex() and (result.imag == 0).all():
|
||||
result = result.real
|
||||
return result
|
||||
|
||||
|
||||
@normalizer
|
||||
@linalg_errors
|
||||
def eigvalsh(a: ArrayLike, UPLO="L"):
|
||||
a = _atleast_float_1(a)
|
||||
return torch.linalg.eigvalsh(a, UPLO=UPLO)
|
||||
@@ -0,0 +1,191 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
"""Wrapper to mimic (parts of) np.random API surface.
|
||||
|
||||
NumPy has strict guarantees on reproducibility etc; here we don't give any.
|
||||
|
||||
Q: default dtype is float64 in numpy
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from math import sqrt
|
||||
|
||||
import torch
|
||||
|
||||
from . import _dtypes_impl, _util
|
||||
from ._normalizations import array_or_scalar, ArrayLike, normalizer
|
||||
|
||||
|
||||
__all__ = [
|
||||
"seed",
|
||||
"random_sample",
|
||||
"sample",
|
||||
"random",
|
||||
"rand",
|
||||
"randn",
|
||||
"normal",
|
||||
"choice",
|
||||
"randint",
|
||||
"shuffle",
|
||||
"uniform",
|
||||
]
|
||||
|
||||
|
||||
def use_numpy_random():
|
||||
# local import to avoid ref cycles
|
||||
import torch._dynamo.config as config
|
||||
|
||||
return config.use_numpy_random_stream
|
||||
|
||||
|
||||
def deco_stream(func):
|
||||
@functools.wraps(func)
|
||||
def inner(*args, **kwds):
|
||||
if not use_numpy_random():
|
||||
return func(*args, **kwds)
|
||||
else:
|
||||
import numpy
|
||||
|
||||
from ._ndarray import ndarray
|
||||
|
||||
f = getattr(numpy.random, func.__name__)
|
||||
|
||||
# numpy funcs accept numpy ndarrays, unwrap
|
||||
args = tuple(
|
||||
arg.tensor.numpy() if isinstance(arg, ndarray) else arg for arg in args
|
||||
)
|
||||
kwds = {
|
||||
key: val.tensor.numpy() if isinstance(val, ndarray) else val
|
||||
for key, val in kwds.items()
|
||||
}
|
||||
|
||||
value = f(*args, **kwds)
|
||||
|
||||
# `value` can be either numpy.ndarray or python scalar (or None)
|
||||
if isinstance(value, numpy.ndarray):
|
||||
value = ndarray(torch.as_tensor(value))
|
||||
|
||||
return value
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
@deco_stream
|
||||
def seed(seed=None):
|
||||
if seed is not None:
|
||||
torch.random.manual_seed(seed)
|
||||
|
||||
|
||||
@deco_stream
|
||||
def random_sample(size=None):
|
||||
if size is None:
|
||||
size = ()
|
||||
dtype = _dtypes_impl.default_dtypes().float_dtype
|
||||
values = torch.empty(size, dtype=dtype).uniform_()
|
||||
return array_or_scalar(values, return_scalar=size == ())
|
||||
|
||||
|
||||
def rand(*size):
|
||||
if size == ():
|
||||
size = None
|
||||
return random_sample(size)
|
||||
|
||||
|
||||
sample = random_sample
|
||||
random = random_sample
|
||||
|
||||
|
||||
@deco_stream
|
||||
def uniform(low=0.0, high=1.0, size=None):
|
||||
if size is None:
|
||||
size = ()
|
||||
dtype = _dtypes_impl.default_dtypes().float_dtype
|
||||
values = torch.empty(size, dtype=dtype).uniform_(low, high)
|
||||
return array_or_scalar(values, return_scalar=size == ())
|
||||
|
||||
|
||||
@deco_stream
|
||||
def randn(*size):
|
||||
dtype = _dtypes_impl.default_dtypes().float_dtype
|
||||
values = torch.randn(size, dtype=dtype)
|
||||
return array_or_scalar(values, return_scalar=size == ())
|
||||
|
||||
|
||||
@deco_stream
|
||||
def normal(loc=0.0, scale=1.0, size=None):
|
||||
if size is None:
|
||||
size = ()
|
||||
dtype = _dtypes_impl.default_dtypes().float_dtype
|
||||
values = torch.empty(size, dtype=dtype).normal_(loc, scale)
|
||||
return array_or_scalar(values, return_scalar=size == ())
|
||||
|
||||
|
||||
@deco_stream
|
||||
def shuffle(x):
|
||||
# no @normalizer because we do not cast e.g. lists to tensors
|
||||
from ._ndarray import ndarray
|
||||
|
||||
if isinstance(x, torch.Tensor):
|
||||
tensor = x
|
||||
elif isinstance(x, ndarray):
|
||||
tensor = x.tensor
|
||||
else:
|
||||
raise NotImplementedError("We do not random.shuffle lists in-place")
|
||||
|
||||
perm = torch.randperm(tensor.shape[0])
|
||||
xp = tensor[perm]
|
||||
tensor.copy_(xp)
|
||||
|
||||
|
||||
@deco_stream
|
||||
def randint(low, high=None, size=None):
|
||||
if size is None:
|
||||
size = ()
|
||||
if not isinstance(size, (tuple, list)):
|
||||
size = (size,)
|
||||
if high is None:
|
||||
low, high = 0, low
|
||||
values = torch.randint(low, high, size=size)
|
||||
return array_or_scalar(values, int, return_scalar=size == ())
|
||||
|
||||
|
||||
@deco_stream
|
||||
@normalizer
|
||||
def choice(a: ArrayLike, size=None, replace=True, p: ArrayLike | None = None):
|
||||
# https://stackoverflow.com/questions/59461811/random-choice-with-pytorch
|
||||
if a.numel() == 1:
|
||||
a = torch.arange(a)
|
||||
|
||||
# TODO: check a.dtype is integer -- cf np.random.choice(3.4) which raises
|
||||
|
||||
# number of draws
|
||||
if size is None:
|
||||
num_el = 1
|
||||
elif _util.is_sequence(size):
|
||||
num_el = 1
|
||||
for el in size:
|
||||
num_el *= el
|
||||
else:
|
||||
num_el = size
|
||||
|
||||
# prepare the probabilities
|
||||
if p is None:
|
||||
p = torch.ones_like(a) / a.shape[0]
|
||||
|
||||
# cf https://github.com/numpy/numpy/blob/main/numpy/random/mtrand.pyx#L973
|
||||
atol = sqrt(torch.finfo(p.dtype).eps)
|
||||
if abs(p.sum() - 1.0) > atol:
|
||||
raise ValueError("probabilities do not sum to 1.")
|
||||
|
||||
# actually sample
|
||||
indices = torch.multinomial(p, num_el, replacement=replace)
|
||||
|
||||
if _util.is_sequence(size):
|
||||
indices = indices.reshape(size)
|
||||
|
||||
samples = a[indices]
|
||||
|
||||
return samples
|
||||
@@ -0,0 +1,20 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from .utils import (
|
||||
_gen_alignment_data,
|
||||
assert_,
|
||||
assert_allclose,
|
||||
assert_almost_equal,
|
||||
assert_array_almost_equal,
|
||||
assert_array_equal,
|
||||
assert_array_less,
|
||||
assert_equal,
|
||||
assert_raises_regex,
|
||||
assert_warns,
|
||||
HAS_REFCOUNT,
|
||||
IS_WASM,
|
||||
suppress_warnings,
|
||||
)
|
||||
|
||||
|
||||
# from .testing import assert_allclose # FIXME
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user