bbqddt2/Antigravity
0
1# mypy: ignore-errors
2
3"""Implementation of reduction operations, to be wrapped into arrays, dtypes etc
4in the 'public' layer.
5
6Anything here only deals with torch objects, e.g. "dtype" is a torch.dtype instance etc
7"""
8
9from __future__ import annotations
10
11import functools
12from typing import TYPE_CHECKING
13
14import torch
15
16from . import _dtypes_impl, _util
17
18
19if TYPE_CHECKING:
20 from ._normalizations import (
21 ArrayLike,
22 AxisLike,
23 DTypeLike,
24 KeepDims,
25 NotImplementedType,
26 OutArray,
27 )
28
29
30def _deco_axis_expand(func):
31 """
32 Generically handle axis arguments in reductions.
33 axis is *always* the 2nd arg in the function so no need to have a look at its signature
34 """
35
36 @functools.wraps(func)
37 def wrapped(a, axis=None, *args, **kwds):
38 if axis is not None:
39 axis = _util.normalize_axis_tuple(axis, a.ndim)
40
41 if axis == ():
42 # So we insert a length-one axis and run the reduction along it.
43 # We cannot return a.clone() as this would sidestep the checks inside the function
44 newshape = _util.expand_shape(a.shape, axis=0)
45 a = a.reshape(newshape)
46 axis = (0,)
47
48 return func(a, axis, *args, **kwds)
49
50 return wrapped
51
52
53def _atleast_float(dtype, other_dtype):
54 """Return a dtype that is real or complex floating-point.
55
56 For inputs that are boolean or integer dtypes, this returns the default
57 float dtype; inputs that are complex get converted to the default complex
58 dtype; real floating-point dtypes (`float*`) get passed through unchanged
59 """
60 if dtype is None:
61 dtype = other_dtype
62 if not (dtype.is_floating_point or dtype.is_complex):
63 return _dtypes_impl.default_dtypes().float_dtype
64 return dtype
65
66
67@_deco_axis_expand
68def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims: KeepDims = False):
69 return a.count_nonzero(axis)
70
71
72@_deco_axis_expand
73def argmax(
74 a: ArrayLike,
75 axis: AxisLike = None,
76 out: OutArray | None = None,
77 *,
78 keepdims: KeepDims = False,
79):
80 if a.is_complex():
81 raise NotImplementedError(f"argmax with dtype={a.dtype}.")
82
83 axis = _util.allow_only_single_axis(axis)
84
85 if a.dtype == torch.bool:
86 # RuntimeError: "argmax_cpu" not implemented for 'Bool'
87 a = a.to(torch.uint8)
88
89 return torch.argmax(a, axis)
90
91
92@_deco_axis_expand
93def argmin(
94 a: ArrayLike,
95 axis: AxisLike = None,
96 out: OutArray | None = None,
97 *,
98 keepdims: KeepDims = False,
99):
100 if a.is_complex():
101 raise NotImplementedError(f"argmin with dtype={a.dtype}.")
102
103 axis = _util.allow_only_single_axis(axis)
104
105 if a.dtype == torch.bool:
106 # RuntimeError: "argmin_cpu" not implemented for 'Bool'
107 a = a.to(torch.uint8)
108
109 return torch.argmin(a, axis)
110
111
112@_deco_axis_expand
113def any(
114 a: ArrayLike,
115 axis: AxisLike = None,
116 out: OutArray | None = None,
117 keepdims: KeepDims = False,
118 *,
119 where: NotImplementedType = None,
120):
121 axis = _util.allow_only_single_axis(axis)
122 axis_kw = {} if axis is None else {"dim": axis}
123 return torch.any(a, **axis_kw)
124
125
126@_deco_axis_expand
127def all(
128 a: ArrayLike,
129 axis: AxisLike = None,
130 out: OutArray | None = None,
131 keepdims: KeepDims = False,
132 *,
133 where: NotImplementedType = None,
134):
135 axis = _util.allow_only_single_axis(axis)
136 axis_kw = {} if axis is None else {"dim": axis}
137 return torch.all(a, **axis_kw)
138
139
140@_deco_axis_expand
141def amax(
142 a: ArrayLike,
143 axis: AxisLike = None,
144 out: OutArray | None = None,
145 keepdims: KeepDims = False,
146 initial: NotImplementedType = None,
147 where: NotImplementedType = None,
148):
149 if a.is_complex():
150 raise NotImplementedError(f"amax with dtype={a.dtype}")
151
152 return a.amax(axis)
153
154
155max = amax
156
157
158@_deco_axis_expand
159def amin(
160 a: ArrayLike,
161 axis: AxisLike = None,
162 out: OutArray | None = None,
163 keepdims: KeepDims = False,
164 initial: NotImplementedType = None,
165 where: NotImplementedType = None,
166):
167 if a.is_complex():
168 raise NotImplementedError(f"amin with dtype={a.dtype}")
169
170 return a.amin(axis)
171
172
173min = amin
174
175
176@_deco_axis_expand
177def ptp(
178 a: ArrayLike,
179 axis: AxisLike = None,
180 out: OutArray | None = None,
181 keepdims: KeepDims = False,
182):
183 return a.amax(axis) - a.amin(axis)
184
185
186@_deco_axis_expand
187def sum(
188 a: ArrayLike,
189 axis: AxisLike = None,
190 dtype: DTypeLike | None = None,
191 out: OutArray | None = None,
192 keepdims: KeepDims = False,
193 initial: NotImplementedType = None,
194 where: NotImplementedType = None,
195):
196 if dtype is not None and not isinstance(dtype, torch.dtype):
197 raise AssertionError(
198 f"dtype must be None or a torch.dtype, got {type(dtype).__name__}"
199 )
200
201 if dtype == torch.bool:
202 dtype = _dtypes_impl.default_dtypes().int_dtype
203
204 axis_kw = {} if axis is None else {"dim": axis}
205 return a.sum(dtype=dtype, **axis_kw)
206
207
208@_deco_axis_expand
209def prod(
210 a: ArrayLike,
211 axis: AxisLike = None,
212 dtype: DTypeLike | None = None,
213 out: OutArray | None = None,
214 keepdims: KeepDims = False,
215 initial: NotImplementedType = None,
216 where: NotImplementedType = None,
217):
218 axis = _util.allow_only_single_axis(axis)
219
220 if dtype == torch.bool:
221 dtype = _dtypes_impl.default_dtypes().int_dtype
222
223 axis_kw = {} if axis is None else {"dim": axis}
224 return a.prod(dtype=dtype, **axis_kw)
225
226
227product = prod
228
229
230@_deco_axis_expand
231def mean(
232 a: ArrayLike,
233 axis: AxisLike = None,
234 dtype: DTypeLike | None = None,
235 out: OutArray | None = None,
236 keepdims: KeepDims = False,
237 *,
238 where: NotImplementedType = None,
239):
240 dtype = _atleast_float(dtype, a.dtype)
241
242 axis_kw = {} if axis is None else {"dim": axis}
243 result = a.mean(dtype=dtype, **axis_kw)
244
245 return result
246
247
248@_deco_axis_expand
249def std(
250 a: ArrayLike,
251 axis: AxisLike = None,
252 dtype: DTypeLike | None = None,
253 out: OutArray | None = None,
254 ddof=0,
255 keepdims: KeepDims = False,
256 *,
257 where: NotImplementedType = None,
258):
259 in_dtype = dtype
260 dtype = _atleast_float(dtype, a.dtype)
261 tensor = _util.cast_if_needed(a, dtype)
262 result = tensor.std(dim=axis, correction=ddof)
263 return _util.cast_if_needed(result, in_dtype)
264
265
266@_deco_axis_expand
267def var(
268 a: ArrayLike,
269 axis: AxisLike = None,
270 dtype: DTypeLike | None = None,
271 out: OutArray | None = None,
272 ddof=0,
273 keepdims: KeepDims = False,
274 *,
275 where: NotImplementedType = None,
276):
277 in_dtype = dtype
278 dtype = _atleast_float(dtype, a.dtype)
279 tensor = _util.cast_if_needed(a, dtype)
280 result = tensor.var(dim=axis, correction=ddof)
281 return _util.cast_if_needed(result, in_dtype)
282
283
284# cumsum / cumprod are almost reductions:
285# 1. no keepdims
286# 2. axis=None flattens
287
288
289def cumsum(
290 a: ArrayLike,
291 axis: AxisLike = None,
292 dtype: DTypeLike | None = None,
293 out: OutArray | None = None,
294):
295 if dtype == torch.bool:
296 dtype = _dtypes_impl.default_dtypes().int_dtype
297 if dtype is None:
298 dtype = a.dtype
299
300 (a,), axis = _util.axis_none_flatten(a, axis=axis)
301 axis = _util.normalize_axis_index(axis, a.ndim)
302
303 return a.cumsum(axis=axis, dtype=dtype)
304
305
306def cumprod(
307 a: ArrayLike,
308 axis: AxisLike = None,
309 dtype: DTypeLike | None = None,
310 out: OutArray | None = None,
311):
312 if dtype == torch.bool:
313 dtype = _dtypes_impl.default_dtypes().int_dtype
314 if dtype is None:
315 dtype = a.dtype
316
317 (a,), axis = _util.axis_none_flatten(a, axis=axis)
318 axis = _util.normalize_axis_index(axis, a.ndim)
319
320 return a.cumprod(axis=axis, dtype=dtype)
321
322
323cumproduct = cumprod
324
325
326def average(
327 a: ArrayLike,
328 axis=None,
329 weights: ArrayLike = None,
330 returned=False,
331 *,
332 keepdims=False,
333):
334 if weights is None:
335 result = mean(a, axis=axis)
336 wsum = torch.as_tensor(a.numel() / result.numel(), dtype=result.dtype)
337 else:
338 if not a.dtype.is_floating_point:
339 a = a.double()
340
341 # axis & weights
342 if a.shape != weights.shape:
343 if axis is None:
344 raise TypeError(
345 "Axis must be specified when shapes of a and weights differ."
346 )
347 if weights.ndim != 1:
348 raise TypeError(
349 "1D weights expected when shapes of a and weights differ."
350 )
351 if weights.shape[0] != a.shape[axis]:
352 raise ValueError(
353 "Length of weights not compatible with specified axis."
354 )
355
356 # setup weight to broadcast along axis
357 weights = torch.broadcast_to(weights, (a.ndim - 1) * (1,) + weights.shape)
358 weights = weights.swapaxes(-1, axis)
359
360 # do the work
361 result_dtype = _dtypes_impl.result_type_impl(a, weights)
362 numerator = sum(a * weights, axis, dtype=result_dtype)
363 wsum = sum(weights, axis, dtype=result_dtype)
364 result = numerator / wsum
365
366 # We process keepdims manually because the decorator does not deal with variadic returns
367 if keepdims:
368 result = _util.apply_keepdims(result, axis, a.ndim)
369
370 if returned:
371 if wsum.shape != result.shape:
372 wsum = torch.broadcast_to(wsum, result.shape).clone()
373 return result, wsum
374 else:
375 return result
376
377
378# Not using deco_axis_expand as it assumes that axis is the second arg
379def quantile(
380 a: ArrayLike,
381 q: ArrayLike,
382 axis: AxisLike = None,
383 out: OutArray | None = None,
384 overwrite_input=False,
385 method="linear",
386 keepdims: KeepDims = False,
387 *,
388 interpolation: NotImplementedType = None,
389):
390 if overwrite_input:
391 # raise NotImplementedError("overwrite_input in quantile not implemented.")
392 # NumPy documents that `overwrite_input` MAY modify inputs:
393 # https://numpy.org/doc/stable/reference/generated/numpy.percentile.html#numpy-percentile
394 # Here we choose to work out-of-place because why not.
395 pass
396
397 if not a.dtype.is_floating_point:
398 dtype = _dtypes_impl.default_dtypes().float_dtype
399 a = a.to(dtype)
400
401 # edge case: torch.quantile only supports float32 and float64
402 if a.dtype == torch.float16:
403 a = a.to(torch.float32)
404
405 if axis is None:
406 a = a.flatten()
407 q = q.flatten()
408 axis = (0,)
409 else:
410 axis = _util.normalize_axis_tuple(axis, a.ndim)
411
412 # FIXME(Mario) Doesn't np.quantile accept a tuple?
413 # torch.quantile does accept a number. If we don't want to implement the tuple behaviour
414 # (it's deffo low prio) change `normalize_axis_tuple` into a normalize_axis index above.
415 axis = _util.allow_only_single_axis(axis)
416
417 q = _util.cast_if_needed(q, a.dtype)
418
419 return torch.quantile(a, q, axis=axis, interpolation=method)
420
421
422def percentile(
423 a: ArrayLike,
424 q: ArrayLike,
425 axis: AxisLike = None,
426 out: OutArray | None = None,
427 overwrite_input=False,
428 method="linear",
429 keepdims: KeepDims = False,
430 *,
431 interpolation: NotImplementedType = None,
432):
433 # np.percentile(float_tensor, 30) : q.dtype is int64 => q / 100.0 is float32
434 if _dtypes_impl.python_type_for_torch(q.dtype) is int:
435 q = q.to(_dtypes_impl.default_dtypes().float_dtype)
436 qq = q / 100.0
437
438 return quantile(
439 a,
440 qq,
441 axis=axis,
442 overwrite_input=overwrite_input,
443 method=method,
444 keepdims=keepdims,
445 interpolation=interpolation,
446 )
447
448
449def median(
450 a: ArrayLike,
451 axis=None,
452 out: OutArray | None = None,
453 overwrite_input=False,
454 keepdims: KeepDims = False,
455):
456 return quantile(
457 a,
458 torch.as_tensor(0.5),
459 axis=axis,
460 overwrite_input=overwrite_input,
461 out=out,
462 keepdims=keepdims,
463 )
464