bbqddt2/Antigravity

1# mypy: ignore-errors
2
3""" "Normalize" arguments: convert array_likes to tensors, dtypes to torch dtypes and so on."""
4
5from __future__ import annotations
6
7import functools
8import inspect
9import operator
10import types
11import typing
12
13import torch
14
15from . import _dtypes, _dtypes_impl, _util
16
17
18ArrayLike = typing.TypeVar("ArrayLike")
19Scalar = int | float | complex | bool
20ArrayLikeOrScalar = ArrayLike | Scalar
21
22DTypeLike = typing.TypeVar("DTypeLike")
23AxisLike = typing.TypeVar("AxisLike")
24NDArray = typing.TypeVar("NDArray")
25CastingModes = typing.TypeVar("CastingModes")
26KeepDims = typing.TypeVar("KeepDims")
27
28# OutArray is to annotate the out= array argument.
29#
30# This one is special is several respects:
31# First, It needs to be an NDArray, and we need to preserve the `result is out`
32# semantics. Therefore, we cannot just extract the Tensor from the out array.
33# So we never pass the out array to implementer functions and handle it in the
34# `normalizer` below.
35# Second, the out= argument can be either keyword or positional argument, and
36# as a positional arg, it can be anywhere in the signature.
37# To handle all this, we define a special `OutArray` annotation and dispatch on it.
38#
39OutArray = typing.TypeVar("OutArray")
40
41NotImplementedType = typing.TypeVar(
42    "NotImplementedType", bound=types.NotImplementedType
43)
44
45
46def normalize_array_like(x, parm=None):  # codespell:ignore
47    from ._ndarray import asarray
48
49    return asarray(x).tensor
50
51
52def normalize_array_like_or_scalar(x, parm=None):  # codespell:ignore
53    if _dtypes_impl.is_scalar_or_symbolic(x):
54        return x
55    return normalize_array_like(x, parm)  # codespell:ignore
56
57
58def normalize_optional_array_like_or_scalar(x, parm=None):  # codespell:ignore
59    if x is None:
60        return None
61    return normalize_array_like_or_scalar(x, parm)  # codespell:ignore
62
63
64def normalize_optional_array_like(x, parm=None):  # codespell:ignore
65    # This explicit normalizer is needed because otherwise normalize_array_like
66    # does not run for a parameter annotated as Optional[ArrayLike]
67    return None if x is None else normalize_array_like(x, parm)  # codespell:ignore
68
69
70def normalize_seq_array_like(x, parm=None):  # codespell:ignore
71    return tuple(normalize_array_like(value) for value in x)
72
73
74def normalize_dtype(dtype, parm=None):  # codespell:ignore
75    # cf _decorators.dtype_to_torch
76    torch_dtype = None
77    if dtype is not None:
78        dtype = _dtypes.dtype(dtype)
79        torch_dtype = dtype.torch_dtype
80    return torch_dtype
81
82
83def normalize_not_implemented(arg, parm):  # codespell:ignore
84    if arg != parm.default:  # codespell:ignore
85        raise NotImplementedError(
86            f"'{parm.name}' parameter is not supported."  # codespell:ignore
87        )
88
89
90def normalize_axis_like(arg, parm=None):  # codespell:ignore
91    from ._ndarray import ndarray
92
93    if isinstance(arg, ndarray):
94        arg = operator.index(arg)
95    return arg
96
97
98def normalize_ndarray(arg, parm=None):  # codespell:ignore
99    # check the arg is an ndarray, extract its tensor attribute
100    if arg is None:
101        return arg
102
103    from ._ndarray import ndarray
104
105    if not isinstance(arg, ndarray):
106        raise TypeError(f"'{parm.name}' must be an array")  # codespell:ignore
107    return arg.tensor
108
109
110def normalize_outarray(arg, parm=None):  # codespell:ignore
111    # almost normalize_ndarray, only return the array, not its tensor
112    if arg is None:
113        return arg
114    from ._ndarray import ndarray
115
116    # Dynamo can pass torch tensors as out arguments,
117    # wrap it in an ndarray before processing
118    if isinstance(arg, torch.Tensor):
119        arg = ndarray(arg)
120
121    if not isinstance(arg, ndarray):
122        raise TypeError(f"'{parm.name}' must be an array")  # codespell:ignore
123    return arg
124
125
126def normalize_casting(arg, parm=None):  # codespell:ignore
127    if arg not in ["no", "equiv", "safe", "same_kind", "unsafe"]:
128        raise ValueError(
129            f"casting must be one of 'no', 'equiv', 'safe', 'same_kind', or 'unsafe' (got '{arg}')"
130        )
131    return arg
132
133
134normalizers = {
135    "ArrayLike": normalize_array_like,
136    "ArrayLikeOrScalar": normalize_array_like_or_scalar,
137    "Optional[ArrayLike]": normalize_optional_array_like,
138    "ArrayLike | None": normalize_optional_array_like,
139    "Sequence[ArrayLike]": normalize_seq_array_like,
140    "Optional[ArrayLikeOrScalar]": normalize_optional_array_like_or_scalar,
141    "ArrayLikeOrScalar | None": normalize_optional_array_like_or_scalar,
142    "Optional[NDArray]": normalize_ndarray,
143    "NDArray | None": normalize_ndarray,
144    "Optional[OutArray]": normalize_outarray,
145    "OutArray | None": normalize_outarray,
146    "NDArray": normalize_ndarray,
147    "Optional[DTypeLike]": normalize_dtype,
148    "DTypeLike | None": normalize_dtype,
149    "AxisLike": normalize_axis_like,
150    "NotImplementedType": normalize_not_implemented,
151    "Optional[CastingModes]": normalize_casting,
152    "CastingModes | None": normalize_casting,
153}
154
155
156def maybe_normalize(arg, parm):  # codespell:ignore
157    """Normalize arg if a normalizer is registered."""
158    normalizer = normalizers.get(parm.annotation)  # codespell:ignore
159    return normalizer(arg, parm) if normalizer else arg  # codespell:ignore
160
161
162# ### Return value helpers ###
163
164
165def maybe_copy_to(out, result, promote_scalar_result=False):
166    # NB: here out is either an ndarray or None
167    if out is None:
168        return result
169    elif isinstance(result, torch.Tensor):
170        if result.shape != out.shape:
171            can_fit = result.numel() == 1 and out.ndim == 0
172            if promote_scalar_result and can_fit:
173                result = result.squeeze()
174            else:
175                raise ValueError(
176                    f"Bad size of the out array: out.shape = {out.shape}"
177                    f" while result.shape = {result.shape}."
178                )
179        out.tensor.copy_(result)
180        return out
181    elif isinstance(result, (tuple, list)):
182        return type(result)(
183            maybe_copy_to(o, r, promote_scalar_result) for o, r in zip(out, result)
184        )
185    else:
186        raise AssertionError  # We should never hit this path
187
188
189def wrap_tensors(result):
190    from ._ndarray import ndarray
191
192    if isinstance(result, torch.Tensor):
193        return ndarray(result)
194    elif isinstance(result, (tuple, list)):
195        result = type(result)(wrap_tensors(x) for x in result)
196    return result
197
198
199def array_or_scalar(values, py_type=float, return_scalar=False):
200    if return_scalar:
201        return py_type(values.item())
202    else:
203        from ._ndarray import ndarray
204
205        return ndarray(values)
206
207
208# ### The main decorator to normalize arguments / postprocess the output ###
209
210
211def normalizer(_func=None, *, promote_scalar_result=False):
212    def normalizer_inner(func):
213        @functools.wraps(func)
214        def wrapped(*args, **kwds):
215            sig = inspect.signature(func)
216            params = sig.parameters
217            first_param = next(iter(params.values()))
218
219            # NumPy's API does not have positional args before variadic positional args
220            if first_param.kind == inspect.Parameter.VAR_POSITIONAL:
221                args = [maybe_normalize(arg, first_param) for arg in args]
222            else:
223                # NB: extra unknown arguments: pass through, will raise in func(*args) below
224                args = (
225                    tuple(
226                        maybe_normalize(arg, parm)  # codespell:ignore
227                        for arg, parm in zip(args, params.values())  # codespell:ignore
228                    )
229                    + args[len(params.values()) :]
230                )
231
232            kwds = {
233                name: maybe_normalize(arg, params[name]) if name in params else arg
234                for name, arg in kwds.items()
235            }
236
237            result = func(*args, **kwds)
238
239            # keepdims
240            bound_args = None
241            if "keepdims" in params and params["keepdims"].annotation == "KeepDims":
242                # keepdims can be in any position so we need sig.bind
243                bound_args = sig.bind(*args, **kwds).arguments
244                if bound_args.get("keepdims", False):
245                    # In this case the first arg is the initial tensor and
246                    # the second arg is (optionally) the axis
247                    tensor = args[0]
248                    axis = bound_args.get("axis")
249                    result = _util.apply_keepdims(result, axis, tensor.ndim)
250
251            # out
252            if "out" in params:
253                # out can be in any position so we need sig.bind
254                if bound_args is None:
255                    bound_args = sig.bind(*args, **kwds).arguments
256                out = bound_args.get("out")
257                result = maybe_copy_to(out, result, promote_scalar_result)
258            result = wrap_tensors(result)
259
260            return result
261
262        return wrapped
263
264    if _func is None:
265        return normalizer_inner
266    else:
267        return normalizer_inner(_func)
268