bbqddt2/Antigravity
0
1# mypy: ignore-errors
2
3"""A thin pytorch / numpy compat layer.
4
5Things imported from here have numpy-compatible signatures but operate on
6pytorch tensors.
7"""
8
9# Contents of this module ends up in the main namespace via _funcs.py
10# where type annotations are used in conjunction with the @normalizer decorator.
11from __future__ import annotations
12
13import builtins
14import itertools
15import operator
16from typing import TYPE_CHECKING
17
18import torch
19
20from . import _dtypes_impl, _util
21
22
23if TYPE_CHECKING:
24 from collections.abc import Sequence
25
26 from ._normalizations import (
27 ArrayLike,
28 ArrayLikeOrScalar,
29 CastingModes,
30 DTypeLike,
31 NDArray,
32 NotImplementedType,
33 OutArray,
34 )
35
36
37def copy(
38 a: ArrayLike, order: NotImplementedType = "K", subok: NotImplementedType = False
39):
40 return a.clone()
41
42
43def copyto(
44 dst: NDArray,
45 src: ArrayLike,
46 casting: CastingModes | None = "same_kind",
47 where: NotImplementedType = None,
48):
49 (src,) = _util.typecast_tensors((src,), dst.dtype, casting=casting)
50 dst.copy_(src)
51
52
53def atleast_1d(*arys: ArrayLike):
54 res = torch.atleast_1d(*arys)
55 if isinstance(res, tuple):
56 return list(res)
57 else:
58 return res
59
60
61def atleast_2d(*arys: ArrayLike):
62 res = torch.atleast_2d(*arys)
63 if isinstance(res, tuple):
64 return list(res)
65 else:
66 return res
67
68
69def atleast_3d(*arys: ArrayLike):
70 res = torch.atleast_3d(*arys)
71 if isinstance(res, tuple):
72 return list(res)
73 else:
74 return res
75
76
77def _concat_check(tup, dtype, out):
78 if tup == ():
79 raise ValueError("need at least one array to concatenate")
80
81 """Check inputs in concatenate et al."""
82 if out is not None and dtype is not None:
83 # mimic numpy
84 raise TypeError(
85 "concatenate() only takes `out` or `dtype` as an "
86 "argument, but both were provided."
87 )
88
89
90def _concat_cast_helper(tensors, out=None, dtype=None, casting="same_kind"):
91 """Figure out dtypes, cast if necessary."""
92
93 if out is not None or dtype is not None:
94 # figure out the type of the inputs and outputs
95 out_dtype = out.dtype.torch_dtype if dtype is None else dtype
96 else:
97 out_dtype = _dtypes_impl.result_type_impl(*tensors)
98
99 # cast input arrays if necessary; do not broadcast them against `out`
100 tensors = _util.typecast_tensors(tensors, out_dtype, casting)
101
102 return tensors
103
104
105def _concatenate(
106 tensors, axis=0, out=None, dtype=None, casting: CastingModes | None = "same_kind"
107):
108 # pure torch implementation, used below and in cov/corrcoef below
109 tensors, axis = _util.axis_none_flatten(*tensors, axis=axis)
110 tensors = _concat_cast_helper(tensors, out, dtype, casting)
111 return torch.cat(tensors, axis)
112
113
114def concatenate(
115 ar_tuple: Sequence[ArrayLike],
116 axis=0,
117 out: OutArray | None = None,
118 dtype: DTypeLike | None = None,
119 casting: CastingModes | None = "same_kind",
120):
121 _concat_check(ar_tuple, dtype, out=out)
122 result = _concatenate(ar_tuple, axis=axis, out=out, dtype=dtype, casting=casting)
123 return result
124
125
126def vstack(
127 tup: Sequence[ArrayLike],
128 *,
129 dtype: DTypeLike | None = None,
130 casting: CastingModes | None = "same_kind",
131):
132 _concat_check(tup, dtype, out=None)
133 tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
134 return torch.vstack(tensors)
135
136
137row_stack = vstack
138
139
140def hstack(
141 tup: Sequence[ArrayLike],
142 *,
143 dtype: DTypeLike | None = None,
144 casting: CastingModes | None = "same_kind",
145):
146 _concat_check(tup, dtype, out=None)
147 tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
148 return torch.hstack(tensors)
149
150
151def dstack(
152 tup: Sequence[ArrayLike],
153 *,
154 dtype: DTypeLike | None = None,
155 casting: CastingModes | None = "same_kind",
156):
157 # XXX: in numpy 1.24 dstack does not have dtype and casting keywords
158 # but {h,v}stack do. Hence add them here for consistency.
159 _concat_check(tup, dtype, out=None)
160 tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
161 return torch.dstack(tensors)
162
163
164def column_stack(
165 tup: Sequence[ArrayLike],
166 *,
167 dtype: DTypeLike | None = None,
168 casting: CastingModes | None = "same_kind",
169):
170 # XXX: in numpy 1.24 column_stack does not have dtype and casting keywords
171 # but row_stack does. (because row_stack is an alias for vstack, really).
172 # Hence add these keywords here for consistency.
173 _concat_check(tup, dtype, out=None)
174 tensors = _concat_cast_helper(tup, dtype=dtype, casting=casting)
175 return torch.column_stack(tensors)
176
177
178def stack(
179 arrays: Sequence[ArrayLike],
180 axis=0,
181 out: OutArray | None = None,
182 *,
183 dtype: DTypeLike | None = None,
184 casting: CastingModes | None = "same_kind",
185):
186 _concat_check(arrays, dtype, out=out)
187
188 tensors = _concat_cast_helper(arrays, dtype=dtype, casting=casting)
189 result_ndim = tensors[0].ndim + 1
190 axis = _util.normalize_axis_index(axis, result_ndim)
191 return torch.stack(tensors, axis=axis)
192
193
194def append(arr: ArrayLike, values: ArrayLike, axis=None):
195 if axis is None:
196 if arr.ndim != 1:
197 arr = arr.flatten()
198 values = values.flatten()
199 axis = arr.ndim - 1
200 return _concatenate((arr, values), axis=axis)
201
202
203# ### split ###
204
205
206def _split_helper(tensor, indices_or_sections, axis, strict=False):
207 if isinstance(indices_or_sections, int):
208 return _split_helper_int(tensor, indices_or_sections, axis, strict)
209 elif isinstance(indices_or_sections, (list, tuple)):
210 # NB: drop split=..., it only applies to split_helper_int
211 return _split_helper_list(tensor, list(indices_or_sections), axis)
212 else:
213 raise TypeError(f"split_helper: {type(indices_or_sections)}")
214
215
216def _split_helper_int(tensor, indices_or_sections, axis, strict=False):
217 if not isinstance(indices_or_sections, int):
218 raise NotImplementedError("split: indices_or_sections")
219
220 axis = _util.normalize_axis_index(axis, tensor.ndim)
221
222 # numpy: l%n chunks of size (l//n + 1), the rest are sized l//n
223 l, n = tensor.shape[axis], indices_or_sections
224
225 if n <= 0:
226 raise ValueError
227
228 if l % n == 0:
229 num, sz = n, l // n
230 lst = [sz] * num
231 else:
232 if strict:
233 raise ValueError("array split does not result in an equal division")
234
235 num, sz = l % n, l // n + 1
236 lst = [sz] * num
237
238 lst += [sz - 1] * (n - num)
239
240 return torch.split(tensor, lst, axis)
241
242
243def _split_helper_list(tensor, indices_or_sections, axis):
244 if not isinstance(indices_or_sections, list):
245 raise NotImplementedError("split: indices_or_sections: list")
246 # numpy expects indices, while torch expects lengths of sections
247 # also, numpy appends zero-size arrays for indices above the shape[axis]
248 lst = [x for x in indices_or_sections if x <= tensor.shape[axis]]
249 num_extra = len(indices_or_sections) - len(lst)
250
251 lst.append(tensor.shape[axis])
252 lst = [
253 lst[0],
254 ] + [a - b for a, b in zip(lst[1:], lst[:-1])]
255 lst += [0] * num_extra
256
257 return torch.split(tensor, lst, axis)
258
259
260def array_split(ary: ArrayLike, indices_or_sections, axis=0):
261 return _split_helper(ary, indices_or_sections, axis)
262
263
264def split(ary: ArrayLike, indices_or_sections, axis=0):
265 return _split_helper(ary, indices_or_sections, axis, strict=True)
266
267
268def hsplit(ary: ArrayLike, indices_or_sections):
269 if ary.ndim == 0:
270 raise ValueError("hsplit only works on arrays of 1 or more dimensions")
271 axis = 1 if ary.ndim > 1 else 0
272 return _split_helper(ary, indices_or_sections, axis, strict=True)
273
274
275def vsplit(ary: ArrayLike, indices_or_sections):
276 if ary.ndim < 2:
277 raise ValueError("vsplit only works on arrays of 2 or more dimensions")
278 return _split_helper(ary, indices_or_sections, 0, strict=True)
279
280
281def dsplit(ary: ArrayLike, indices_or_sections):
282 if ary.ndim < 3:
283 raise ValueError("dsplit only works on arrays of 3 or more dimensions")
284 return _split_helper(ary, indices_or_sections, 2, strict=True)
285
286
287def kron(a: ArrayLike, b: ArrayLike):
288 return torch.kron(a, b)
289
290
291def vander(x: ArrayLike, N=None, increasing=False):
292 return torch.vander(x, N, increasing)
293
294
295# ### linspace, geomspace, logspace and arange ###
296
297
298def linspace(
299 start: ArrayLike,
300 stop: ArrayLike,
301 num=50,
302 endpoint=True,
303 retstep=False,
304 dtype: DTypeLike | None = None,
305 axis=0,
306):
307 if axis != 0 or retstep or not endpoint:
308 raise NotImplementedError
309 if dtype is None:
310 dtype = _dtypes_impl.default_dtypes().float_dtype
311 # XXX: raises TypeError if start or stop are not scalars
312 return torch.linspace(start, stop, num, dtype=dtype)
313
314
315def geomspace(
316 start: ArrayLike,
317 stop: ArrayLike,
318 num=50,
319 endpoint=True,
320 dtype: DTypeLike | None = None,
321 axis=0,
322):
323 if axis != 0 or not endpoint:
324 raise NotImplementedError
325 base = torch.pow(stop / start, 1.0 / (num - 1))
326 logbase = torch.log(base)
327 return torch.logspace(
328 torch.log(start) / logbase,
329 torch.log(stop) / logbase,
330 num,
331 base=base,
332 )
333
334
335def logspace(
336 start,
337 stop,
338 num=50,
339 endpoint=True,
340 base=10.0,
341 dtype: DTypeLike | None = None,
342 axis=0,
343):
344 if axis != 0 or not endpoint:
345 raise NotImplementedError
346 return torch.logspace(start, stop, num, base=base, dtype=dtype)
347
348
349def arange(
350 start: ArrayLikeOrScalar | None = None,
351 stop: ArrayLikeOrScalar | None = None,
352 step: ArrayLikeOrScalar | None = 1,
353 dtype: DTypeLike | None = None,
354 *,
355 like: NotImplementedType = None,
356):
357 if step == 0:
358 raise ZeroDivisionError
359 if stop is None and start is None:
360 raise TypeError
361 if stop is None:
362 # XXX: this breaks if start is passed as a kwarg:
363 # arange(start=4) should raise (no stop) but doesn't
364 start, stop = 0, start
365 if start is None:
366 start = 0
367
368 # the dtype of the result
369 if dtype is None:
370 dtype = (
371 _dtypes_impl.default_dtypes().float_dtype
372 if any(_dtypes_impl.is_float_or_fp_tensor(x) for x in (start, stop, step))
373 else _dtypes_impl.default_dtypes().int_dtype
374 )
375 work_dtype = torch.float64 if dtype.is_complex else dtype
376
377 # RuntimeError: "lt_cpu" not implemented for 'ComplexFloat'. Fall back to eager.
378 if any(_dtypes_impl.is_complex_or_complex_tensor(x) for x in (start, stop, step)):
379 raise NotImplementedError
380
381 if (step > 0 and start > stop) or (step < 0 and start < stop):
382 # empty range
383 return torch.empty(0, dtype=dtype)
384
385 result = torch.arange(start, stop, step, dtype=work_dtype)
386 result = _util.cast_if_needed(result, dtype)
387 return result
388
389
390# ### zeros/ones/empty/full ###
391
392
393def empty(
394 shape,
395 dtype: DTypeLike | None = None,
396 order: NotImplementedType = "C",
397 *,
398 like: NotImplementedType = None,
399):
400 if dtype is None:
401 dtype = _dtypes_impl.default_dtypes().float_dtype
402 return torch.empty(shape, dtype=dtype)
403
404
405# NB: *_like functions deliberately deviate from numpy: it has subok=True
406# as the default; we set subok=False and raise on anything else.
407
408
409def empty_like(
410 prototype: ArrayLike,
411 dtype: DTypeLike | None = None,
412 order: NotImplementedType = "K",
413 subok: NotImplementedType = False,
414 shape=None,
415):
416 result = torch.empty_like(prototype, dtype=dtype)
417 if shape is not None:
418 result = result.reshape(shape)
419 return result
420
421
422def full(
423 shape,
424 fill_value: ArrayLike,
425 dtype: DTypeLike | None = None,
426 order: NotImplementedType = "C",
427 *,
428 like: NotImplementedType = None,
429):
430 if isinstance(shape, int):
431 shape = (shape,)
432 if dtype is None:
433 dtype = fill_value.dtype
434 if not isinstance(shape, (tuple, list)):
435 shape = (shape,)
436 return torch.full(shape, fill_value, dtype=dtype)
437
438
439def full_like(
440 a: ArrayLike,
441 fill_value,
442 dtype: DTypeLike | None = None,
443 order: NotImplementedType = "K",
444 subok: NotImplementedType = False,
445 shape=None,
446):
447 # XXX: fill_value broadcasts
448 result = torch.full_like(a, fill_value, dtype=dtype)
449 if shape is not None:
450 result = result.reshape(shape)
451 return result
452
453
454def ones(
455 shape,
456 dtype: DTypeLike | None = None,
457 order: NotImplementedType = "C",
458 *,
459 like: NotImplementedType = None,
460):
461 if dtype is None:
462 dtype = _dtypes_impl.default_dtypes().float_dtype
463 return torch.ones(shape, dtype=dtype)
464
465
466def ones_like(
467 a: ArrayLike,
468 dtype: DTypeLike | None = None,
469 order: NotImplementedType = "K",
470 subok: NotImplementedType = False,
471 shape=None,
472):
473 result = torch.ones_like(a, dtype=dtype)
474 if shape is not None:
475 result = result.reshape(shape)
476 return result
477
478
479def zeros(
480 shape,
481 dtype: DTypeLike | None = None,
482 order: NotImplementedType = "C",
483 *,
484 like: NotImplementedType = None,
485):
486 if dtype is None:
487 dtype = _dtypes_impl.default_dtypes().float_dtype
488 return torch.zeros(shape, dtype=dtype)
489
490
491def zeros_like(
492 a: ArrayLike,
493 dtype: DTypeLike | None = None,
494 order: NotImplementedType = "K",
495 subok: NotImplementedType = False,
496 shape=None,
497):
498 result = torch.zeros_like(a, dtype=dtype)
499 if shape is not None:
500 result = result.reshape(shape)
501 return result
502
503
504# ### cov & corrcoef ###
505
506
507def _xy_helper_corrcoef(x_tensor, y_tensor=None, rowvar=True):
508 """Prepare inputs for cov and corrcoef."""
509
510 # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/function_base.py#L2636
511 if y_tensor is not None:
512 # make sure x and y are at least 2D
513 ndim_extra = 2 - x_tensor.ndim
514 if ndim_extra > 0:
515 x_tensor = x_tensor.view((1,) * ndim_extra + x_tensor.shape)
516 if not rowvar and x_tensor.shape[0] != 1:
517 x_tensor = x_tensor.mT
518 x_tensor = x_tensor.clone()
519
520 ndim_extra = 2 - y_tensor.ndim
521 if ndim_extra > 0:
522 y_tensor = y_tensor.view((1,) * ndim_extra + y_tensor.shape)
523 if not rowvar and y_tensor.shape[0] != 1:
524 y_tensor = y_tensor.mT
525 y_tensor = y_tensor.clone()
526
527 x_tensor = _concatenate((x_tensor, y_tensor), axis=0)
528
529 return x_tensor
530
531
532def corrcoef(
533 x: ArrayLike,
534 y: ArrayLike | None = None,
535 rowvar=True,
536 bias=None,
537 ddof=None,
538 *,
539 dtype: DTypeLike | None = None,
540):
541 if bias is not None or ddof is not None:
542 # deprecated in NumPy
543 raise NotImplementedError
544 xy_tensor = _xy_helper_corrcoef(x, y, rowvar)
545
546 is_half = (xy_tensor.dtype == torch.float16) and xy_tensor.is_cpu
547 if is_half:
548 # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
549 dtype = torch.float32
550
551 xy_tensor = _util.cast_if_needed(xy_tensor, dtype)
552 result = torch.corrcoef(xy_tensor)
553
554 if is_half:
555 result = result.to(torch.float16)
556
557 return result
558
559
560def cov(
561 m: ArrayLike,
562 y: ArrayLike | None = None,
563 rowvar=True,
564 bias=False,
565 ddof=None,
566 fweights: ArrayLike | None = None,
567 aweights: ArrayLike | None = None,
568 *,
569 dtype: DTypeLike | None = None,
570):
571 m = _xy_helper_corrcoef(m, y, rowvar)
572
573 if ddof is None:
574 ddof = 1 if bias == 0 else 0
575
576 is_half = (m.dtype == torch.float16) and m.is_cpu
577 if is_half:
578 # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
579 dtype = torch.float32
580
581 m = _util.cast_if_needed(m, dtype)
582 result = torch.cov(m, correction=ddof, aweights=aweights, fweights=fweights)
583
584 if is_half:
585 result = result.to(torch.float16)
586
587 return result
588
589
590def _conv_corr_impl(a, v, mode):
591 dt = _dtypes_impl.result_type_impl(a, v)
592 a = _util.cast_if_needed(a, dt)
593 v = _util.cast_if_needed(v, dt)
594
595 padding = v.shape[0] - 1 if mode == "full" else mode
596
597 if padding == "same" and v.shape[0] % 2 == 0:
598 # UserWarning: Using padding='same' with even kernel lengths and odd
599 # dilation may require a zero-padded copy of the input be created
600 # (Triggered internally at pytorch/aten/src/ATen/native/Convolution.cpp:1010.)
601 raise NotImplementedError("mode='same' and even-length weights")
602
603 # NumPy only accepts 1D arrays; PyTorch requires 2D inputs and 3D weights
604 aa = a[None, :]
605 vv = v[None, None, :]
606
607 result = torch.nn.functional.conv1d(aa, vv, padding=padding)
608
609 # torch returns a 2D result, numpy returns a 1D array
610 return result[0, :]
611
612
613def convolve(a: ArrayLike, v: ArrayLike, mode="full"):
614 # NumPy: if v is longer than a, the arrays are swapped before computation
615 if a.shape[0] < v.shape[0]:
616 a, v = v, a
617
618 # flip the weights since numpy does and torch does not
619 v = torch.flip(v, (0,))
620
621 return _conv_corr_impl(a, v, mode)
622
623
624def correlate(a: ArrayLike, v: ArrayLike, mode="valid"):
625 v = torch.conj_physical(v)
626 return _conv_corr_impl(a, v, mode)
627
628
629# ### logic & element selection ###
630
631
632def bincount(x: ArrayLike, /, weights: ArrayLike | None = None, minlength=0):
633 if x.numel() == 0:
634 # edge case allowed by numpy
635 x = x.new_empty(0, dtype=int)
636
637 int_dtype = _dtypes_impl.default_dtypes().int_dtype
638 (x,) = _util.typecast_tensors((x,), int_dtype, casting="safe")
639
640 return torch.bincount(x, weights, minlength)
641
642
643def where(
644 condition: ArrayLike,
645 x: ArrayLikeOrScalar | None = None,
646 y: ArrayLikeOrScalar | None = None,
647 /,
648):
649 if (x is None) != (y is None):
650 raise ValueError("either both or neither of x and y should be given")
651
652 if condition.dtype != torch.bool:
653 condition = condition.to(torch.bool)
654
655 if x is None and y is None:
656 result = torch.where(condition)
657 else:
658 result = torch.where(condition, x, y)
659 return result
660
661
662# ###### module-level queries of object properties
663
664
665def ndim(a: ArrayLike):
666 return a.ndim
667
668
669def shape(a: ArrayLike):
670 return tuple(a.shape)
671
672
673def size(a: ArrayLike, axis=None):
674 if axis is None:
675 return a.numel()
676 else:
677 return a.shape[axis]
678
679
680# ###### shape manipulations and indexing
681
682
683def expand_dims(a: ArrayLike, axis):
684 shape = _util.expand_shape(a.shape, axis)
685 return a.view(shape) # never copies
686
687
688def flip(m: ArrayLike, axis=None):
689 # XXX: semantic difference: np.flip returns a view, torch.flip copies
690 if axis is None:
691 axis = tuple(range(m.ndim))
692 else:
693 axis = _util.normalize_axis_tuple(axis, m.ndim)
694 return torch.flip(m, axis)
695
696
697def flipud(m: ArrayLike):
698 return torch.flipud(m)
699
700
701def fliplr(m: ArrayLike):
702 return torch.fliplr(m)
703
704
705def rot90(m: ArrayLike, k=1, axes=(0, 1)):
706 axes = _util.normalize_axis_tuple(axes, m.ndim)
707 return torch.rot90(m, k, axes)
708
709
710# ### broadcasting and indices ###
711
712
713def broadcast_to(array: ArrayLike, shape, subok: NotImplementedType = False):
714 return torch.broadcast_to(array, size=shape)
715
716
717# This is a function from tuples to tuples, so we just reuse it. However,
718# dynamo expects its __module__ to be torch._numpy
719def broadcast_shapes(*args):
720 return torch.broadcast_shapes(*args)
721
722
723def broadcast_arrays(*args: ArrayLike, subok: NotImplementedType = False):
724 return torch.broadcast_tensors(*args)
725
726
727def meshgrid(*xi: ArrayLike, copy=True, sparse=False, indexing="xy"):
728 ndim = len(xi)
729
730 if indexing not in ["xy", "ij"]:
731 raise ValueError("Valid values for `indexing` are 'xy' and 'ij'.")
732
733 s0 = (1,) * ndim
734 output = [x.reshape(s0[:i] + (-1,) + s0[i + 1 :]) for i, x in enumerate(xi)]
735
736 if indexing == "xy" and ndim > 1:
737 # switch first and second axis
738 output[0] = output[0].reshape((1, -1) + s0[2:])
739 output[1] = output[1].reshape((-1, 1) + s0[2:])
740
741 if not sparse:
742 # Return the full N-D matrix (not only the 1-D vector)
743 output = torch.broadcast_tensors(*output)
744
745 if copy:
746 output = [x.clone() for x in output]
747
748 return list(output) # match numpy, return a list
749
750
751def indices(dimensions, dtype: DTypeLike | None = int, sparse=False):
752 # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1691-L1791
753 dimensions = tuple(dimensions)
754 N = len(dimensions)
755 shape = (1,) * N
756 if sparse:
757 res = ()
758 else:
759 res = torch.empty((N,) + dimensions, dtype=dtype)
760 for i, dim in enumerate(dimensions):
761 idx = torch.arange(dim, dtype=dtype).reshape(
762 shape[:i] + (dim,) + shape[i + 1 :]
763 )
764 if sparse:
765 res = res + (idx,)
766 else:
767 res[i] = idx
768 return res
769
770
771# ### tri*-something ###
772
773
774def tril(m: ArrayLike, k=0):
775 return torch.tril(m, k)
776
777
778def triu(m: ArrayLike, k=0):
779 return torch.triu(m, k)
780
781
782def tril_indices(n, k=0, m=None):
783 if m is None:
784 m = n
785 return torch.tril_indices(n, m, offset=k)
786
787
788def triu_indices(n, k=0, m=None):
789 if m is None:
790 m = n
791 return torch.triu_indices(n, m, offset=k)
792
793
794def tril_indices_from(arr: ArrayLike, k=0):
795 if arr.ndim != 2:
796 raise ValueError("input array must be 2-d")
797 # Return a tensor rather than a tuple to avoid a graphbreak
798 return torch.tril_indices(arr.shape[0], arr.shape[1], offset=k)
799
800
801def triu_indices_from(arr: ArrayLike, k=0):
802 if arr.ndim != 2:
803 raise ValueError("input array must be 2-d")
804 # Return a tensor rather than a tuple to avoid a graphbreak
805 return torch.triu_indices(arr.shape[0], arr.shape[1], offset=k)
806
807
808def tri(
809 N,
810 M=None,
811 k=0,
812 dtype: DTypeLike | None = None,
813 *,
814 like: NotImplementedType = None,
815):
816 if M is None:
817 M = N
818 tensor = torch.ones((N, M), dtype=dtype)
819 return torch.tril(tensor, diagonal=k)
820
821
822# ### equality, equivalence, allclose ###
823
824
825def isclose(a: ArrayLike, b: ArrayLike, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
826 dtype = _dtypes_impl.result_type_impl(a, b)
827 a = _util.cast_if_needed(a, dtype)
828 b = _util.cast_if_needed(b, dtype)
829 return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
830
831
832def allclose(a: ArrayLike, b: ArrayLike, rtol=1e-05, atol=1e-08, equal_nan=False):
833 dtype = _dtypes_impl.result_type_impl(a, b)
834 a = _util.cast_if_needed(a, dtype)
835 b = _util.cast_if_needed(b, dtype)
836 return torch.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
837
838
839def _tensor_equal(a1, a2, equal_nan=False):
840 # Implementation of array_equal/array_equiv.
841 if a1.shape != a2.shape:
842 return False
843 cond = a1 == a2
844 if equal_nan:
845 cond = cond | (torch.isnan(a1) & torch.isnan(a2))
846 return cond.all().item()
847
848
849def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan=False):
850 return _tensor_equal(a1, a2, equal_nan=equal_nan)
851
852
853def array_equiv(a1: ArrayLike, a2: ArrayLike):
854 # *almost* the same as array_equal: _equiv tries to broadcast, _equal does not
855 try:
856 a1_t, a2_t = torch.broadcast_tensors(a1, a2)
857 except RuntimeError:
858 # failed to broadcast => not equivalent
859 return False
860 return _tensor_equal(a1_t, a2_t)
861
862
863def nan_to_num(
864 x: ArrayLike, copy: NotImplementedType = True, nan=0.0, posinf=None, neginf=None
865):
866 # work around RuntimeError: "nan_to_num" not implemented for 'ComplexDouble'
867 if x.is_complex():
868 re = torch.nan_to_num(x.real, nan=nan, posinf=posinf, neginf=neginf)
869 im = torch.nan_to_num(x.imag, nan=nan, posinf=posinf, neginf=neginf)
870 return re + 1j * im
871 else:
872 return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
873
874
875# ### put/take_along_axis ###
876
877
878def take(
879 a: ArrayLike,
880 indices: ArrayLike,
881 axis=None,
882 out: OutArray | None = None,
883 mode: NotImplementedType = "raise",
884):
885 (a,), axis = _util.axis_none_flatten(a, axis=axis)
886 axis = _util.normalize_axis_index(axis, a.ndim)
887 idx = (slice(None),) * axis + (indices, ...)
888 result = a[idx]
889 return result
890
891
892def take_along_axis(arr: ArrayLike, indices: ArrayLike, axis):
893 (arr,), axis = _util.axis_none_flatten(arr, axis=axis)
894 axis = _util.normalize_axis_index(axis, arr.ndim)
895 return torch.take_along_dim(arr, indices, axis)
896
897
898def put(
899 a: NDArray,
900 indices: ArrayLike,
901 values: ArrayLike,
902 mode: NotImplementedType = "raise",
903):
904 v = values.type(a.dtype)
905 # If indices is larger than v, expand v to at least the size of indices. Any
906 # unnecessary trailing elements are then trimmed.
907 if indices.numel() > v.numel():
908 ratio = (indices.numel() + v.numel() - 1) // v.numel()
909 v = v.unsqueeze(0).expand((ratio,) + v.shape)
910 # Trim unnecessary elements, regardless if v was expanded or not. Note
911 # np.put() trims v to match indices by default too.
912 if indices.numel() < v.numel():
913 v = v.flatten()
914 v = v[: indices.numel()]
915 a.put_(indices, v)
916 return None
917
918
919def put_along_axis(arr: ArrayLike, indices: ArrayLike, values: ArrayLike, axis):
920 (arr,), axis = _util.axis_none_flatten(arr, axis=axis)
921 axis = _util.normalize_axis_index(axis, arr.ndim)
922
923 indices, values = torch.broadcast_tensors(indices, values)
924 values = _util.cast_if_needed(values, arr.dtype)
925 result = torch.scatter(arr, axis, indices, values)
926 arr.copy_(result.reshape(arr.shape))
927 return None
928
929
930def choose(
931 a: ArrayLike,
932 choices: Sequence[ArrayLike],
933 out: OutArray | None = None,
934 mode: NotImplementedType = "raise",
935):
936 # First, broadcast elements of `choices`
937 choices = torch.stack(torch.broadcast_tensors(*choices))
938
939 # Use an analog of `gather(choices, 0, a)` which broadcasts `choices` vs `a`:
940 # (taken from https://github.com/pytorch/pytorch/issues/9407#issuecomment-1427907939)
941 idx_list = [
942 torch.arange(dim).view((1,) * i + (dim,) + (1,) * (choices.ndim - i - 1))
943 for i, dim in enumerate(choices.shape)
944 ]
945
946 idx_list[0] = a
947 return choices[tuple(idx_list)].squeeze(0)
948
949
950# ### unique et al. ###
951
952
953def unique(
954 ar: ArrayLike,
955 return_index: NotImplementedType = False,
956 return_inverse=False,
957 return_counts=False,
958 axis=None,
959 *,
960 equal_nan: NotImplementedType = True,
961):
962 (ar,), axis = _util.axis_none_flatten(ar, axis=axis)
963 axis = _util.normalize_axis_index(axis, ar.ndim)
964
965 result = torch.unique(
966 ar, return_inverse=return_inverse, return_counts=return_counts, dim=axis
967 )
968
969 return result
970
971
972def nonzero(a: ArrayLike):
973 return torch.nonzero(a, as_tuple=True)
974
975
976def argwhere(a: ArrayLike):
977 return torch.argwhere(a)
978
979
980def flatnonzero(a: ArrayLike):
981 return torch.flatten(a).nonzero(as_tuple=True)[0]
982
983
984def clip(
985 a: ArrayLike,
986 min: ArrayLike | None = None,
987 max: ArrayLike | None = None,
988 out: OutArray | None = None,
989):
990 return torch.clamp(a, min, max)
991
992
993def repeat(a: ArrayLike, repeats: ArrayLikeOrScalar, axis=None):
994 return torch.repeat_interleave(a, repeats, axis)
995
996
997def tile(A: ArrayLike, reps):
998 if isinstance(reps, int):
999 reps = (reps,)
1000 return torch.tile(A, reps)
1001
1002
1003def resize(a: ArrayLike, new_shape=None):
1004 # implementation vendored from
1005 # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/fromnumeric.py#L1420-L1497
1006 if new_shape is None:
1007 return a
1008
1009 if isinstance(new_shape, int):
1010 new_shape = (new_shape,)
1011
1012 a = a.flatten()
1013
1014 new_size = 1
1015 for dim_length in new_shape:
1016 new_size *= dim_length
1017 if dim_length < 0:
1018 raise ValueError("all elements of `new_shape` must be non-negative")
1019
1020 if a.numel() == 0 or new_size == 0:
1021 # First case must zero fill. The second would have repeats == 0.
1022 return torch.zeros(new_shape, dtype=a.dtype)
1023
1024 repeats = -(-new_size // a.numel()) # ceil division
1025 a = concatenate((a,) * repeats)[:new_size]
1026
1027 return reshape(a, new_shape)
1028
1029
1030# ### diag et al. ###
1031
1032
1033def diagonal(a: ArrayLike, offset=0, axis1=0, axis2=1):
1034 axis1 = _util.normalize_axis_index(axis1, a.ndim)
1035 axis2 = _util.normalize_axis_index(axis2, a.ndim)
1036 return torch.diagonal(a, offset, axis1, axis2)
1037
1038
1039def trace(
1040 a: ArrayLike,
1041 offset=0,
1042 axis1=0,
1043 axis2=1,
1044 dtype: DTypeLike | None = None,
1045 out: OutArray | None = None,
1046):
1047 result = torch.diagonal(a, offset, dim1=axis1, dim2=axis2).sum(-1, dtype=dtype)
1048 return result
1049
1050
1051def eye(
1052 N,
1053 M=None,
1054 k=0,
1055 dtype: DTypeLike | None = None,
1056 order: NotImplementedType = "C",
1057 *,
1058 like: NotImplementedType = None,
1059):
1060 if dtype is None:
1061 dtype = _dtypes_impl.default_dtypes().float_dtype
1062 if M is None:
1063 M = N
1064 z = torch.zeros(N, M, dtype=dtype)
1065 z.diagonal(k).fill_(1)
1066 return z
1067
1068
1069def identity(n, dtype: DTypeLike | None = None, *, like: NotImplementedType = None):
1070 return torch.eye(n, dtype=dtype)
1071
1072
1073def diag(v: ArrayLike, k=0):
1074 return torch.diag(v, k)
1075
1076
1077def diagflat(v: ArrayLike, k=0):
1078 return torch.diagflat(v, k)
1079
1080
1081def diag_indices(n, ndim=2):
1082 idx = torch.arange(n)
1083 return (idx,) * ndim
1084
1085
1086def diag_indices_from(arr: ArrayLike):
1087 if not arr.ndim >= 2:
1088 raise ValueError("input array must be at least 2-d")
1089 # For more than d=2, the strided formula is only valid for arrays with
1090 # all dimensions equal, so we check first.
1091 s = arr.shape
1092 if s[1:] != s[:-1]:
1093 raise ValueError("All dimensions of input must be of equal length")
1094 return diag_indices(s[0], arr.ndim)
1095
1096
1097def fill_diagonal(a: ArrayLike, val: ArrayLike, wrap=False):
1098 if a.ndim < 2:
1099 raise ValueError("array must be at least 2-d")
1100 if val.numel() == 0 and not wrap:
1101 a.fill_diagonal_(val)
1102 return a
1103
1104 if val.ndim == 0:
1105 val = val.unsqueeze(0)
1106
1107 # torch.Tensor.fill_diagonal_ only accepts scalars
1108 # If the size of val is too large, then val is trimmed
1109 if a.ndim == 2:
1110 tall = a.shape[0] > a.shape[1]
1111 # wrap does nothing for wide matrices...
1112 if not wrap or not tall:
1113 # Never wraps
1114 diag = a.diagonal()
1115 diag.copy_(val[: diag.numel()])
1116 else:
1117 # wraps and tall... leaving one empty line between diagonals?!
1118 max_, min_ = a.shape
1119 idx = torch.arange(max_ - max_ // (min_ + 1))
1120 mod = idx % min_
1121 div = idx // min_
1122 a[(div * (min_ + 1) + mod, mod)] = val[: idx.numel()]
1123 else:
1124 idx = diag_indices_from(a)
1125 # a.shape = (n, n, ..., n)
1126 a[idx] = val[: a.shape[0]]
1127
1128 return a
1129
1130
1131def vdot(a: ArrayLike, b: ArrayLike, /):
1132 # 1. torch only accepts 1D arrays, numpy flattens
1133 # 2. torch requires matching dtype, while numpy casts (?)
1134 t_a, t_b = torch.atleast_1d(a, b)
1135 if t_a.ndim > 1:
1136 t_a = t_a.flatten()
1137 if t_b.ndim > 1:
1138 t_b = t_b.flatten()
1139
1140 dtype = _dtypes_impl.result_type_impl(t_a, t_b)
1141 is_half = dtype == torch.float16 and (t_a.is_cpu or t_b.is_cpu)
1142 is_bool = dtype == torch.bool
1143
1144 # work around torch's "dot" not implemented for 'Half', 'Bool'
1145 if is_half:
1146 dtype = torch.float32
1147 elif is_bool:
1148 dtype = torch.uint8
1149
1150 t_a = _util.cast_if_needed(t_a, dtype)
1151 t_b = _util.cast_if_needed(t_b, dtype)
1152
1153 result = torch.vdot(t_a, t_b)
1154
1155 if is_half:
1156 result = result.to(torch.float16)
1157 elif is_bool:
1158 result = result.to(torch.bool)
1159
1160 return result
1161
1162
1163def tensordot(a: ArrayLike, b: ArrayLike, axes=2):
1164 if isinstance(axes, (list, tuple)):
1165 axes = [[ax] if isinstance(ax, int) else ax for ax in axes]
1166
1167 target_dtype = _dtypes_impl.result_type_impl(a, b)
1168 a = _util.cast_if_needed(a, target_dtype)
1169 b = _util.cast_if_needed(b, target_dtype)
1170
1171 return torch.tensordot(a, b, dims=axes)
1172
1173
1174def dot(a: ArrayLike, b: ArrayLike, out: OutArray | None = None):
1175 dtype = _dtypes_impl.result_type_impl(a, b)
1176 is_bool = dtype == torch.bool
1177 if is_bool:
1178 dtype = torch.uint8
1179
1180 a = _util.cast_if_needed(a, dtype)
1181 b = _util.cast_if_needed(b, dtype)
1182
1183 if a.ndim == 0 or b.ndim == 0:
1184 result = a * b
1185 else:
1186 result = torch.matmul(a, b)
1187
1188 if is_bool:
1189 result = result.to(torch.bool)
1190
1191 return result
1192
1193
1194def inner(a: ArrayLike, b: ArrayLike, /):
1195 dtype = _dtypes_impl.result_type_impl(a, b)
1196 is_half = dtype == torch.float16 and (a.is_cpu or b.is_cpu)
1197 is_bool = dtype == torch.bool
1198
1199 if is_half:
1200 # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"