bbqddt2/Antigravity
0
1# mypy: ignore-errors
2
3from __future__ import annotations
4
5import builtins
6import math
7import operator
8from collections.abc import Sequence
9
10import torch
11
12from . import _dtypes, _dtypes_impl, _funcs, _ufuncs, _util
13from ._normalizations import (
14 ArrayLike,
15 normalize_array_like,
16 normalizer,
17 NotImplementedType,
18)
19
20
21newaxis = None
22
23FLAGS = [
24 "C_CONTIGUOUS",
25 "F_CONTIGUOUS",
26 "OWNDATA",
27 "WRITEABLE",
28 "ALIGNED",
29 "WRITEBACKIFCOPY",
30 "FNC",
31 "FORC",
32 "BEHAVED",
33 "CARRAY",
34 "FARRAY",
35]
36
37SHORTHAND_TO_FLAGS = {
38 "C": "C_CONTIGUOUS",
39 "F": "F_CONTIGUOUS",
40 "O": "OWNDATA",
41 "W": "WRITEABLE",
42 "A": "ALIGNED",
43 "X": "WRITEBACKIFCOPY",
44 "B": "BEHAVED",
45 "CA": "CARRAY",
46 "FA": "FARRAY",
47}
48
49
50class Flags:
51 def __init__(self, flag_to_value: dict):
52 invalid_keys = [k for k in flag_to_value if k not in FLAGS]
53 if invalid_keys:
54 raise AssertionError(f"Invalid flag keys: {invalid_keys}")
55 self._flag_to_value = flag_to_value
56
57 def __getattr__(self, attr: str):
58 if attr.islower() and attr.upper() in FLAGS:
59 return self[attr.upper()]
60 else:
61 raise AttributeError(f"No flag attribute '{attr}'")
62
63 def __getitem__(self, key):
64 if key in SHORTHAND_TO_FLAGS:
65 key = SHORTHAND_TO_FLAGS[key]
66 if key in FLAGS:
67 try:
68 return self._flag_to_value[key]
69 except KeyError as e:
70 raise NotImplementedError(f"{key=}") from e
71 else:
72 raise KeyError(f"No flag key '{key}'")
73
74 def __setattr__(self, attr, value):
75 if attr.islower() and attr.upper() in FLAGS:
76 self[attr.upper()] = value
77 else:
78 super().__setattr__(attr, value)
79
80 def __setitem__(self, key, value):
81 if key in FLAGS or key in SHORTHAND_TO_FLAGS:
82 raise NotImplementedError("Modifying flags is not implemented")
83 else:
84 raise KeyError(f"No flag key '{key}'")
85
86
87def create_method(fn, name=None):
88 name = name or fn.__name__
89
90 def f(*args, **kwargs):
91 return fn(*args, **kwargs)
92
93 f.__name__ = name
94 f.__qualname__ = f"ndarray.{name}"
95 return f
96
97
98# Map ndarray.name_method -> np.name_func
99# If name_func == None, it means that name_method == name_func
100methods = {
101 "clip": None,
102 "nonzero": None,
103 "repeat": None,
104 "round": None,
105 "squeeze": None,
106 "swapaxes": None,
107 "ravel": None,
108 # linalg
109 "diagonal": None,
110 "dot": None,
111 "trace": None,
112 # sorting
113 "argsort": None,
114 "searchsorted": None,
115 # reductions
116 "argmax": None,
117 "argmin": None,
118 "any": None,
119 "all": None,
120 "max": None,
121 "min": None,
122 "ptp": None,
123 "sum": None,
124 "prod": None,
125 "mean": None,
126 "var": None,
127 "std": None,
128 # scans
129 "cumsum": None,
130 "cumprod": None,
131 # advanced indexing
132 "take": None,
133 "choose": None,
134}
135
136dunder = {
137 "abs": "absolute",
138 "invert": None,
139 "pos": "positive",
140 "neg": "negative",
141 "gt": "greater",
142 "lt": "less",
143 "ge": "greater_equal",
144 "le": "less_equal",
145}
146
147# dunder methods with right-looking and in-place variants
148ri_dunder = {
149 "add": None,
150 "sub": "subtract",
151 "mul": "multiply",
152 "truediv": "divide",
153 "floordiv": "floor_divide",
154 "pow": "power",
155 "mod": "remainder",
156 "and": "bitwise_and",
157 "or": "bitwise_or",
158 "xor": "bitwise_xor",
159 "lshift": "left_shift",
160 "rshift": "right_shift",
161 "matmul": None,
162}
163
164
165def _upcast_int_indices(index):
166 if isinstance(index, torch.Tensor):
167 if index.dtype in (torch.int8, torch.int16, torch.int32, torch.uint8):
168 return index.to(torch.int64)
169 elif isinstance(index, tuple):
170 return tuple(_upcast_int_indices(i) for i in index)
171 return index
172
173
174def _has_advanced_indexing(index):
175 """Check if there's any advanced indexing"""
176 return any(
177 isinstance(idx, (Sequence, bool))
178 or (isinstance(idx, torch.Tensor) and (idx.dtype == torch.bool or idx.ndim > 0))
179 for idx in index
180 )
181
182
183def _numpy_compatible_indexing(index):
184 """Convert scalar indices to lists when advanced indexing is present for NumPy compatibility."""
185 if not isinstance(index, tuple):
186 index = (index,)
187
188 # Check if there's any advanced indexing (sequences, booleans, or tensors)
189 has_advanced = _has_advanced_indexing(index)
190
191 if not has_advanced:
192 return index
193
194 # Convert integer scalar indices to single-element lists when advanced indexing is present
195 # Note: Do NOT convert boolean scalars (True/False) as they have special meaning in NumPy
196 converted = []
197 for idx in index:
198 if isinstance(idx, int) and not isinstance(idx, bool):
199 # Integer scalars should be converted to lists
200 converted.append([idx])
201 elif (
202 isinstance(idx, torch.Tensor)
203 and idx.ndim == 0
204 and not torch.is_floating_point(idx)
205 and idx.dtype != torch.bool
206 ):
207 # Zero-dimensional tensors holding integers should be treated the same as integer scalars
208 converted.append([idx])
209 else:
210 # Everything else (booleans, lists, slices, etc.) stays as is
211 converted.append(idx)
212
213 return tuple(converted)
214
215
216def _get_bool_depth(s):
217 """Returns the depth of a boolean sequence/tensor"""
218 if isinstance(s, bool):
219 return True, 0
220 if isinstance(s, torch.Tensor) and s.dtype == torch.bool:
221 return True, s.ndim
222 if not (isinstance(s, Sequence) and s and s[0] != s):
223 return False, 0
224 is_bool, depth = _get_bool_depth(s[0])
225 return is_bool, depth + 1
226
227
228def _numpy_empty_ellipsis_patch(index, tensor_ndim):
229 """
230 Patch for NumPy-compatible ellipsis behavior when ellipsis doesn't match any dimensions.
231
232 In NumPy, when an ellipsis (...) doesn't actually match any dimensions of the input array,
233 it still acts as a separator between advanced indices. PyTorch doesn't have this behavior.
234
235 This function detects when we have:
236 1. Advanced indexing on both sides of an ellipsis
237 2. The ellipsis doesn't actually match any dimensions
238 """
239 if not isinstance(index, tuple):
240 index = (index,)
241
242 # Find ellipsis position
243 ellipsis_pos = None
244 for i, idx in enumerate(index):
245 if idx is Ellipsis:
246 ellipsis_pos = i
247 break
248
249 # If no ellipsis, no patch needed
250 if ellipsis_pos is None:
251 return index, lambda x: x, lambda x: x
252
253 # Count non-ellipsis dimensions consumed by the index
254 consumed_dims = 0
255 for idx in index:
256 is_bool, depth = _get_bool_depth(idx)
257 if is_bool:
258 consumed_dims += depth
259 elif idx is Ellipsis or idx is None:
260 continue
261 else:
262 consumed_dims += 1
263
264 # Calculate how many dimensions the ellipsis should match
265 ellipsis_dims = tensor_ndim - consumed_dims
266
267 # Check if ellipsis doesn't match any dimensions
268 if ellipsis_dims == 0:
269 # Check if we have advanced indexing on both sides of ellipsis
270 left_advanced = _has_advanced_indexing(index[:ellipsis_pos])
271 right_advanced = _has_advanced_indexing(index[ellipsis_pos + 1 :])
272
273 if left_advanced and right_advanced:
274 # This is the case where NumPy and PyTorch differ
275 # We need to ensure the advanced indices are treated as separated
276 new_index = index[:ellipsis_pos] + (None,) + index[ellipsis_pos + 1 :]
277 end_ndims = 1 + sum(
278 1 for idx in index[ellipsis_pos + 1 :] if isinstance(idx, slice)
279 )
280
281 def squeeze_fn(x):
282 return x.squeeze(-end_ndims)
283
284 def unsqueeze_fn(x):
285 if isinstance(x, torch.Tensor) and x.ndim >= end_ndims:
286 return x.unsqueeze(-end_ndims)
287 return x
288
289 return new_index, squeeze_fn, unsqueeze_fn
290
291 return index, lambda x: x, lambda x: x
292
293
294# Used to indicate that a parameter is unspecified (as opposed to explicitly
295# `None`)
296class _Unspecified:
297 pass
298
299
300_Unspecified.unspecified = _Unspecified()
301
302###############################################################
303# ndarray class #
304###############################################################
305
306
307class ndarray:
308 def __init__(self, t=None):
309 if t is None:
310 self.tensor = torch.Tensor()
311 elif isinstance(t, torch.Tensor):
312 self.tensor = t
313 else:
314 raise ValueError(
315 "ndarray constructor is not recommended; prefer"
316 "either array(...) or zeros/empty(...)"
317 )
318
319 # Register NumPy functions as methods
320 for method, name in methods.items():
321 fn = getattr(_funcs, name or method)
322 vars()[method] = create_method(fn, method)
323
324 # Regular methods but coming from ufuncs
325 conj = create_method(_ufuncs.conjugate, "conj")
326 conjugate = create_method(_ufuncs.conjugate)
327
328 for method, name in dunder.items():
329 fn = getattr(_ufuncs, name or method)
330 method = f"__{method}__"
331 vars()[method] = create_method(fn, method)
332
333 for method, name in ri_dunder.items():
334 fn = getattr(_ufuncs, name or method)
335 plain = f"__{method}__"
336 vars()[plain] = create_method(fn, plain)
337 rvar = f"__r{method}__"
338 vars()[rvar] = create_method(lambda self, other, fn=fn: fn(other, self), rvar)
339 ivar = f"__i{method}__"
340 vars()[ivar] = create_method(
341 lambda self, other, fn=fn: fn(self, other, out=self), ivar
342 )
343
344 # There's no __idivmod__
345 __divmod__ = create_method(_ufuncs.divmod, "__divmod__")
346 __rdivmod__ = create_method(
347 lambda self, other: _ufuncs.divmod(other, self), "__rdivmod__"
348 )
349
350 # prevent loop variables leaking into the ndarray class namespace
351 del ivar, rvar, name, plain, fn, method
352
353 @property
354 def shape(self):
355 return tuple(self.tensor.shape)
356
357 @property
358 def size(self):
359 return self.tensor.numel()
360
361 @property
362 def ndim(self):
363 return self.tensor.ndim
364
365 @property
366 def dtype(self):
367 return _dtypes.dtype(self.tensor.dtype)
368
369 @property
370 def strides(self):
371 elsize = self.tensor.element_size()
372 return tuple(stride * elsize for stride in self.tensor.stride())
373
374 @property
375 def itemsize(self):
376 return self.tensor.element_size()
377
378 @property
379 def flags(self):
380 # Note contiguous in torch is assumed C-style
381 return Flags(
382 {
383 "C_CONTIGUOUS": self.tensor.is_contiguous(),
384 "F_CONTIGUOUS": self.T.tensor.is_contiguous(),
385 "OWNDATA": self.tensor._base is None,
386 "WRITEABLE": True, # pytorch does not have readonly tensors
387 }
388 )
389
390 @property
391 def data(self):
392 return self.tensor.data_ptr()
393
394 @property
395 def nbytes(self):
396 return self.tensor.storage().nbytes()
397
398 @property
399 def T(self):
400 return self.transpose()
401
402 @property
403 def real(self):
404 return _funcs.real(self)
405
406 @real.setter
407 def real(self, value):
408 self.tensor.real = asarray(value).tensor
409
410 @property
411 def imag(self):
412 return _funcs.imag(self)
413
414 @imag.setter
415 def imag(self, value):
416 self.tensor.imag = asarray(value).tensor
417
418 @property
419 def flat(self):
420 return self.ravel()
421
422 # ctors
423 def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
424 if order != "K":
425 raise NotImplementedError(f"astype(..., order={order} is not implemented.")
426 if casting != "unsafe":
427 raise NotImplementedError(
428 f"astype(..., casting={casting} is not implemented."
429 )
430 if not subok:
431 raise NotImplementedError(f"astype(..., subok={subok} is not implemented.")
432 if not copy:
433 raise NotImplementedError(f"astype(..., copy={copy} is not implemented.")
434 torch_dtype = _dtypes.dtype(dtype).torch_dtype
435 t = self.tensor.to(torch_dtype)
436 return ndarray(t)
437
438 @normalizer
439 def copy(self: ArrayLike, order: NotImplementedType = "C"):
440 return self.clone()
441
442 @normalizer
443 def flatten(self: ArrayLike, order: NotImplementedType = "C"):
444 return torch.flatten(self)
445
446 def resize(self, *new_shape, refcheck=False):
447 # NB: differs from np.resize: fills with zeros instead of making repeated copies of input.
448 if refcheck:
449 raise NotImplementedError(
450 f"resize(..., refcheck={refcheck} is not implemented."
451 )
452 if new_shape in [(), (None,)]:
453 return
454
455 # support both x.resize((2, 2)) and x.resize(2, 2)
456 if len(new_shape) == 1:
457 new_shape = new_shape[0]
458 if isinstance(new_shape, int):
459 new_shape = (new_shape,)
460
461 if builtins.any(x < 0 for x in new_shape):
462 raise ValueError("all elements of `new_shape` must be non-negative")
463
464 new_numel, old_numel = math.prod(new_shape), self.tensor.numel()
465
466 self.tensor.resize_(new_shape)
467
468 if new_numel >= old_numel:
469 # zero-fill new elements
470 if not self.tensor.is_contiguous():
471 raise AssertionError("tensor must be contiguous for resize with growth")
472 b = self.tensor.flatten() # does not copy
473 b[old_numel:].zero_()
474
475 def view(self, dtype=_Unspecified.unspecified, type=_Unspecified.unspecified):
476 if dtype is _Unspecified.unspecified:
477 dtype = self.dtype
478 if type is not _Unspecified.unspecified:
479 raise NotImplementedError(f"view(..., type={type} is not implemented.")
480 torch_dtype = _dtypes.dtype(dtype).torch_dtype
481 tview = self.tensor.view(torch_dtype)
482 return ndarray(tview)
483
484 @normalizer
485 def fill(self, value: ArrayLike):
486 # Both Pytorch and NumPy accept 0D arrays/tensors and scalars, and
487 # error out on D > 0 arrays
488 self.tensor.fill_(value)
489
490 def tolist(self):
491 return self.tensor.tolist()
492
493 def __iter__(self):
494 return (ndarray(x) for x in self.tensor.__iter__())
495
496 def __str__(self):
497 return (
498 str(self.tensor)
499 .replace("tensor", "torch.ndarray")
500 .replace("dtype=torch.", "dtype=")
501 )
502
503 __repr__ = create_method(__str__)
504
505 def __eq__(self, other):
506 try:
507 return _ufuncs.equal(self, other)
508 except (RuntimeError, TypeError):
509 # Failed to convert other to array: definitely not equal.
510 falsy = torch.full(self.shape, fill_value=False, dtype=bool)
511 return asarray(falsy)
512
513 def __ne__(self, other):
514 return ~(self == other)
515
516 def __index__(self):
517 try:
518 return operator.index(self.tensor.item())
519 except Exception as exc:
520 raise TypeError(
521 "only integer scalar arrays can be converted to a scalar index"
522 ) from exc
523
524 def __bool__(self):
525 return bool(self.tensor)
526
527 def __int__(self):
528 return int(self.tensor)
529
530 def __float__(self):
531 return float(self.tensor)
532
533 def __complex__(self):
534 return complex(self.tensor)
535
536 def is_integer(self):
537 try:
538 v = self.tensor.item()
539 result = int(v) == v
540 except Exception:
541 result = False
542 return result
543
544 def __len__(self):
545 return self.tensor.shape[0]
546
547 def __contains__(self, x):
548 return self.tensor.__contains__(x)
549
550 def transpose(self, *axes):
551 # np.transpose(arr, axis=None) but arr.transpose(*axes)
552 return _funcs.transpose(self, axes)
553
554 def reshape(self, *shape, order="C"):
555 # arr.reshape(shape) and arr.reshape(*shape)
556 return _funcs.reshape(self, shape, order=order)
557
558 def sort(self, axis=-1, kind=None, order=None):
559 # ndarray.sort works in-place
560 _funcs.copyto(self, _funcs.sort(self, axis, kind, order))
561
562 def item(self, *args):
563 # Mimic NumPy's implementation with three special cases (no arguments,
564 # a flat index and a multi-index):
565 # https://github.com/numpy/numpy/blob/main/numpy/_core/src/multiarray/methods.c#L702
566 if args == ():
567 return self.tensor.item()
568 elif len(args) == 1:
569 # int argument
570 return self.ravel()[args[0]]
571 else:
572 return self.__getitem__(args)
573
574 def __getitem__(self, index):
575 tensor = self.tensor
576
577 def neg_step(i, s):
578 if not (isinstance(s, slice) and s.step is not None and s.step < 0):
579 return s
580
581 nonlocal tensor
582 tensor = torch.flip(tensor, (i,))
583
584 # Account for the fact that a slice includes the start but not the end
585 if not (isinstance(s.start, int) or s.start is None):
586 raise AssertionError(
587 f"slice start must be int or None, got {type(s.start).__name__}"
588 )
589 if not (isinstance(s.stop, int) or s.stop is None):
590 raise AssertionError(
591 f"slice stop must be int or None, got {type(s.stop).__name__}"
592 )
593 start = s.stop + 1 if s.stop else None
594 stop = s.start + 1 if s.start else None
595
596 return slice(start, stop, -s.step)
597
598 if isinstance(index, Sequence):
599 index = type(index)(neg_step(i, s) for i, s in enumerate(index))
600 else:
601 index = neg_step(0, index)
602 index = _util.ndarrays_to_tensors(index)
603 index = _upcast_int_indices(index)
604 # Apply NumPy-compatible indexing conversion
605 index = _numpy_compatible_indexing(index)
606 # Apply NumPy-compatible empty ellipsis behavior
607 index, maybe_squeeze, _ = _numpy_empty_ellipsis_patch(index, tensor.ndim)
608 return maybe_squeeze(ndarray(tensor.__getitem__(index)))
609
610 def __setitem__(self, index, value):
611 index = _util.ndarrays_to_tensors(index)
612 index = _upcast_int_indices(index)
613 # Apply NumPy-compatible indexing conversion
614 index = _numpy_compatible_indexing(index)
615 # Apply NumPy-compatible empty ellipsis behavior
616 index, _, maybe_unsqueeze = _numpy_empty_ellipsis_patch(index, self.tensor.ndim)
617
618 if not _dtypes_impl.is_scalar(value):
619 value = normalize_array_like(value)
620 value = _util.cast_if_needed(value, self.tensor.dtype)
621
622 return self.tensor.__setitem__(index, maybe_unsqueeze(value))
623
624 take = _funcs.take
625 put = _funcs.put
626
627 def __dlpack__(self, *, stream=None):
628 return self.tensor.__dlpack__(stream=stream)
629
630 def __dlpack_device__(self):
631 return self.tensor.__dlpack_device__()
632
633
634def _tolist(obj):
635 """Recursively convert tensors into lists."""
636 a1 = []
637 for elem in obj:
638 if isinstance(elem, (list, tuple)):
639 elem = _tolist(elem)
640 if isinstance(elem, ndarray):
641 a1.append(elem.tensor.tolist())
642 else:
643 a1.append(elem)
644 return a1
645
646
647# This is the ideally the only place which talks to ndarray directly.
648# The rest goes through asarray (preferred) or array.
649
650
651def array(obj, dtype=None, *, copy=True, order="K", subok=False, ndmin=0, like=None):
652 if subok is not False:
653 raise NotImplementedError("'subok' parameter is not supported.")
654 if like is not None:
655 raise NotImplementedError("'like' parameter is not supported.")
656 if order != "K":
657 raise NotImplementedError
658
659 # a happy path
660 if (
661 isinstance(obj, ndarray)
662 and copy is False
663 and dtype is None
664 and ndmin <= obj.ndim
665 ):
666 return obj
667
668 if isinstance(obj, (list, tuple)):
669 # FIXME and they have the same dtype, device, etc
670 if obj and all(isinstance(x, torch.Tensor) for x in obj):
671 # list of arrays: *under torch.Dynamo* these are FakeTensors
672 obj = torch.stack(obj)
673 else:
674 # XXX: remove tolist
675 # lists of ndarrays: [1, [2, 3], ndarray(4)] convert to lists of lists
676 obj = _tolist(obj)
677
678 # is obj an ndarray already?
679 if isinstance(obj, ndarray):
680 obj = obj.tensor
681
682 # is a specific dtype requested?
683 torch_dtype = None
684 if dtype is not None:
685 torch_dtype = _dtypes.dtype(dtype).torch_dtype
686
687 tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin)
688 return ndarray(tensor)
689
690
691def asarray(a, dtype=None, order="K", *, like=None):
692 return array(a, dtype=dtype, order=order, like=like, copy=False, ndmin=0)
693
694
695def ascontiguousarray(a, dtype=None, *, like=None):
696 arr = asarray(a, dtype=dtype, like=like)
697 if not arr.tensor.is_contiguous():
698 arr.tensor = arr.tensor.contiguous()
699 return arr
700
701
702def from_dlpack(x, /):
703 t = torch.from_dlpack(x)
704 return ndarray(t)
705
706
707def _extract_dtype(entry):
708 try:
709 dty = _dtypes.dtype(entry)
710 except Exception:
711 dty = asarray(entry).dtype
712 return dty
713
714
715def can_cast(from_, to, casting="safe"):
716 from_ = _extract_dtype(from_)
717 to_ = _extract_dtype(to)
718
719 return _dtypes_impl.can_cast_impl(from_.torch_dtype, to_.torch_dtype, casting)
720
721
722def result_type(*arrays_and_dtypes):
723 tensors = []
724 for entry in arrays_and_dtypes:
725 try:
726 t = asarray(entry).tensor
727 except (RuntimeError, ValueError, TypeError):
728 dty = _dtypes.dtype(entry)
729 t = torch.empty(1, dtype=dty.torch_dtype)
730 tensors.append(t)
731
732 torch_dtype = _dtypes_impl.result_type_impl(*tensors)
733 return _dtypes.dtype(torch_dtype)
734