bbqddt2/Antigravity
0
1# mypy: ignore-errors
2
3"""Assorted utilities, which do not need anything other then torch and stdlib."""
4
5import operator
6
7import torch
8
9from . import _dtypes_impl
10
11
12# https://github.com/numpy/numpy/blob/v1.23.0/numpy/distutils/misc_util.py#L497-L504
13def is_sequence(seq):
14 if isinstance(seq, str):
15 return False
16 try:
17 len(seq)
18 except Exception:
19 return False
20 return True
21
22
23class AxisError(ValueError, IndexError):
24 pass
25
26
27class UFuncTypeError(TypeError, RuntimeError):
28 pass
29
30
31def cast_if_needed(tensor, dtype):
32 # NB: no casting if dtype=None
33 if dtype is not None and tensor.dtype != dtype:
34 tensor = tensor.to(dtype)
35 return tensor
36
37
38def cast_int_to_float(x):
39 # cast integers and bools to the default float dtype
40 if _dtypes_impl._category(x.dtype) < 2:
41 x = x.to(_dtypes_impl.default_dtypes().float_dtype)
42 return x
43
44
45# a replica of the version in ./numpy/numpy/core/src/multiarray/common.h
46def normalize_axis_index(ax, ndim, argname=None):
47 if not (-ndim <= ax < ndim):
48 raise AxisError(f"axis {ax} is out of bounds for array of dimension {ndim}")
49 if ax < 0:
50 ax += ndim
51 return ax
52
53
54# from https://github.com/numpy/numpy/blob/main/numpy/core/numeric.py#L1378
55def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
56 """
57 Normalizes an axis argument into a tuple of non-negative integer axes.
58
59 This handles shorthands such as ``1`` and converts them to ``(1,)``,
60 as well as performing the handling of negative indices covered by
61 `normalize_axis_index`.
62
63 By default, this forbids axes from being specified multiple times.
64 Used internally by multi-axis-checking logic.
65
66 Parameters
67 ----------
68 axis : int, iterable of int
69 The un-normalized index or indices of the axis.
70 ndim : int
71 The number of dimensions of the array that `axis` should be normalized
72 against.
73 argname : str, optional
74 A prefix to put before the error message, typically the name of the
75 argument.
76 allow_duplicate : bool, optional
77 If False, the default, disallow an axis from being specified twice.
78
79 Returns
80 -------
81 normalized_axes : tuple of int
82 The normalized axis index, such that `0 <= normalized_axis < ndim`
83 """
84 # Optimization to speed-up the most common cases.
85 if type(axis) not in (tuple, list):
86 try:
87 axis = [operator.index(axis)]
88 except TypeError:
89 pass
90 # Going via an iterator directly is slower than via list comprehension.
91 axis = tuple(normalize_axis_index(ax, ndim, argname) for ax in axis)
92 if not allow_duplicate and len(set(map(int, axis))) != len(axis):
93 if argname:
94 raise ValueError(f"repeated axis in `{argname}` argument")
95 else:
96 raise ValueError("repeated axis")
97 return axis
98
99
100def allow_only_single_axis(axis):
101 if axis is None:
102 return axis
103 if len(axis) != 1:
104 raise NotImplementedError("does not handle tuple axis")
105 return axis[0]
106
107
108def expand_shape(arr_shape, axis):
109 # taken from numpy 1.23.x, expand_dims function
110 if type(axis) not in (list, tuple):
111 axis = (axis,)
112 out_ndim = len(axis) + len(arr_shape)
113 axis = normalize_axis_tuple(axis, out_ndim)
114 shape_it = iter(arr_shape)
115 shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)]
116 return shape
117
118
119def apply_keepdims(tensor, axis, ndim):
120 if axis is None:
121 # tensor was a scalar
122 shape = (1,) * ndim
123 tensor = tensor.expand(shape).contiguous()
124 else:
125 shape = expand_shape(tensor.shape, axis)
126 tensor = tensor.reshape(shape)
127 return tensor
128
129
130def axis_none_flatten(*tensors, axis=None):
131 """Flatten the arrays if axis is None."""
132 if axis is None:
133 tensors = tuple(ar.flatten() for ar in tensors)
134 return tensors, 0
135 else:
136 return tensors, axis
137
138
139def typecast_tensor(t, target_dtype, casting):
140 """Dtype-cast tensor to target_dtype.
141
142 Parameters
143 ----------
144 t : torch.Tensor
145 The tensor to cast
146 target_dtype : torch dtype object
147 The array dtype to cast all tensors to
148 casting : str
149 The casting mode, see `np.can_cast`
150
151 Returns
152 -------
153 `torch.Tensor` of the `target_dtype` dtype
154
155 Raises
156 ------
157 ValueError
158 if the argument cannot be cast according to the `casting` rule
159
160 """
161 can_cast = _dtypes_impl.can_cast_impl
162
163 if not can_cast(t.dtype, target_dtype, casting=casting):
164 raise TypeError(
165 f"Cannot cast array data from {t.dtype} to"
166 f" {target_dtype} according to the rule '{casting}'"
167 )
168 return cast_if_needed(t, target_dtype)
169
170
171def typecast_tensors(tensors, target_dtype, casting):
172 return tuple(typecast_tensor(t, target_dtype, casting) for t in tensors)
173
174
175def _try_convert_to_tensor(obj):
176 try:
177 tensor = torch.as_tensor(obj)
178 except Exception as e:
179 mesg = f"failed to convert {obj} to ndarray. \nInternal error is: {str(e)}."
180 raise NotImplementedError(mesg) # noqa: B904
181 return tensor
182
183
184def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
185 """The core logic of the array(...) function.
186
187 Parameters
188 ----------
189 obj : tensor_like
190 The thing to coerce
191 dtype : torch.dtype object or None
192 Coerce to this torch dtype
193 copy : bool
194 Copy or not
195 ndmin : int
196 The results as least this many dimensions
197 is_weak : bool
198 Whether obj is a weakly typed python scalar.
199
200 Returns
201 -------
202 tensor : torch.Tensor
203 a tensor object with requested dtype, ndim and copy semantics.
204
205 Notes
206 -----
207 This is almost a "tensor_like" coercive function. Does not handle wrapper
208 ndarrays (those should be handled in the ndarray-aware layer prior to
209 invoking this function).
210 """
211 if isinstance(obj, torch.Tensor):
212 tensor = obj
213 else:
214 # tensor.dtype is the pytorch default, typically float32. If obj's elements
215 # are not exactly representable in float32, we've lost precision:
216 # >>> torch.as_tensor(1e12).item() - 1e12
217 # -4096.0
218 default_dtype = torch.get_default_dtype()
219 torch.set_default_dtype(_dtypes_impl.get_default_dtype_for(torch.float32))
220 try:
221 tensor = _try_convert_to_tensor(obj)
222 finally:
223 torch.set_default_dtype(default_dtype)
224
225 # type cast if requested
226 tensor = cast_if_needed(tensor, dtype)
227
228 # adjust ndim if needed
229 ndim_extra = ndmin - tensor.ndim
230 if ndim_extra > 0:
231 tensor = tensor.view((1,) * ndim_extra + tensor.shape)
232
233 # special handling for np._CopyMode
234 try:
235 copy = bool(copy)
236 except ValueError:
237 # TODO handle _CopyMode.IF_NEEDED correctly
238 copy = False
239 # copy if requested
240 if copy:
241 tensor = tensor.clone()
242
243 return tensor
244
245
246def ndarrays_to_tensors(*inputs):
247 """Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
248 from ._ndarray import ndarray
249
250 if len(inputs) == 0:
251 return ValueError()
252 elif len(inputs) == 1:
253 input_ = inputs[0]
254 if isinstance(input_, ndarray):
255 return input_.tensor
256 elif isinstance(input_, tuple):
257 result = []
258 for sub_input in input_:
259 sub_result = ndarrays_to_tensors(sub_input)
260 result.append(sub_result)
261 return tuple(result)
262 else:
263 return input_
264 else:
265 if not isinstance(inputs, tuple):
266 raise AssertionError(
267 f"Expected inputs to be a tuple, got {type(inputs).__name__}"
268 )
269 return ndarrays_to_tensors(inputs)
270