bbqddt2/Antigravity
0
1# mypy: ignore-errors
2
3"""Define analogs of numpy dtypes supported by pytorch.
4Define the scalar types and supported dtypes and numpy <--> torch dtype mappings.
5"""
6
7import builtins
8
9import torch
10
11from . import _dtypes_impl
12
13
14# ### Scalar types ###
15
16
17class generic:
18 name = "generic"
19
20 def __new__(cls, value):
21 # NumPy scalars are modelled as 0-D arrays
22 # so a call to np.float32(4) produces a 0-D array.
23
24 from ._ndarray import asarray, ndarray
25
26 if isinstance(value, str) and value in ["inf", "nan"]:
27 value = {"inf": torch.inf, "nan": torch.nan}[value]
28
29 if isinstance(value, ndarray):
30 return value.astype(cls)
31 else:
32 return asarray(value, dtype=cls)
33
34
35##################
36# abstract types #
37##################
38
39
40class number(generic):
41 name = "number"
42
43
44class integer(number):
45 name = "integer"
46
47
48class inexact(number):
49 name = "inexact"
50
51
52class signedinteger(integer):
53 name = "signedinteger"
54
55
56class unsignedinteger(integer):
57 name = "unsignedinteger"
58
59
60class floating(inexact):
61 name = "floating"
62
63
64class complexfloating(inexact):
65 name = "complexfloating"
66
67
68_abstract_dtypes = [
69 "generic",
70 "number",
71 "integer",
72 "signedinteger",
73 "unsignedinteger",
74 "inexact",
75 "floating",
76 "complexfloating",
77]
78
79# ##### concrete types
80
81# signed integers
82
83
84class int8(signedinteger):
85 name = "int8"
86 typecode = "b"
87 torch_dtype = torch.int8
88
89
90class int16(signedinteger):
91 name = "int16"
92 typecode = "h"
93 torch_dtype = torch.int16
94
95
96class int32(signedinteger):
97 name = "int32"
98 typecode = "i"
99 torch_dtype = torch.int32
100
101
102class int64(signedinteger):
103 name = "int64"
104 typecode = "l"
105 torch_dtype = torch.int64
106
107
108# unsigned integers
109
110
111class uint8(unsignedinteger):
112 name = "uint8"
113 typecode = "B"
114 torch_dtype = torch.uint8
115
116
117class uint16(unsignedinteger):
118 name = "uint16"
119 typecode = "H"
120 torch_dtype = torch.uint16
121
122
123class uint32(signedinteger):
124 name = "uint32"
125 typecode = "I"
126 torch_dtype = torch.uint32
127
128
129class uint64(signedinteger):
130 name = "uint64"
131 typecode = "L"
132 torch_dtype = torch.uint64
133
134
135# floating point
136
137
138class float16(floating):
139 name = "float16"
140 typecode = "e"
141 torch_dtype = torch.float16
142
143
144class float32(floating):
145 name = "float32"
146 typecode = "f"
147 torch_dtype = torch.float32
148
149
150class float64(floating):
151 name = "float64"
152 typecode = "d"
153 torch_dtype = torch.float64
154
155
156class complex64(complexfloating):
157 name = "complex64"
158 typecode = "F"
159 torch_dtype = torch.complex64
160
161
162class complex128(complexfloating):
163 name = "complex128"
164 typecode = "D"
165 torch_dtype = torch.complex128
166
167
168class bool_(generic):
169 name = "bool_"
170 typecode = "?"
171 torch_dtype = torch.bool
172
173
174# name aliases
175_name_aliases = {
176 "intp": int64,
177 "int_": int64,
178 "intc": int32,
179 "byte": int8,
180 "short": int16,
181 "longlong": int64, # XXX: is this correct?
182 "ulonglong": uint64,
183 "ubyte": uint8,
184 "half": float16,
185 "single": float32,
186 "double": float64,
187 "float_": float64,
188 "csingle": complex64,
189 "singlecomplex": complex64,
190 "cdouble": complex128,
191 "cfloat": complex128,
192 "complex_": complex128,
193}
194# We register float_ = float32 and so on
195for name, obj in _name_aliases.items():
196 vars()[name] = obj
197
198
199# Replicate this NumPy-defined way of grouping scalar types,
200# cf tests/core/test_scalar_methods.py
201sctypes = {
202 "int": [int8, int16, int32, int64],
203 "uint": [uint8, uint16, uint32, uint64],
204 "float": [float16, float32, float64],
205 "complex": [complex64, complex128],
206 "others": [bool_],
207}
208
209
210# Support mappings/functions
211
212_names = {st.name: st for cat in sctypes for st in sctypes[cat]}
213_typecodes = {st.typecode: st for cat in sctypes for st in sctypes[cat]}
214_torch_dtypes = {st.torch_dtype: st for cat in sctypes for st in sctypes[cat]}
215
216
217_aliases = {
218 "u1": uint8,
219 "i1": int8,
220 "i2": int16,
221 "i4": int32,
222 "i8": int64,
223 "b": int8, # XXX: srsly?
224 "f2": float16,
225 "f4": float32,
226 "f8": float64,
227 "c8": complex64,
228 "c16": complex128,
229 # numpy-specific trailing underscore
230 "bool_": bool_,
231}
232
233
234_python_types = {
235 int: int64,
236 float: float64,
237 complex: complex128,
238 builtins.bool: bool_,
239 # also allow stringified names of python types
240 int.__name__: int64,
241 float.__name__: float64,
242 complex.__name__: complex128,
243 builtins.bool.__name__: bool_,
244}
245
246
247def sctype_from_string(s):
248 """Normalize a string value: a type 'name' or a typecode or a width alias."""
249 if s in _names:
250 return _names[s]
251 if s in _name_aliases:
252 return _name_aliases[s]
253 if s in _typecodes:
254 return _typecodes[s]
255 if s in _aliases:
256 return _aliases[s]
257 if s in _python_types:
258 return _python_types[s]
259 raise TypeError(f"data type {s!r} not understood")
260
261
262def sctype_from_torch_dtype(torch_dtype):
263 return _torch_dtypes[torch_dtype]
264
265
266# ### DTypes. ###
267
268
269def dtype(arg):
270 if arg is None:
271 arg = _dtypes_impl.default_dtypes().float_dtype
272 return DType(arg)
273
274
275class DType:
276 def __init__(self, arg):
277 # a pytorch object?
278 if isinstance(arg, torch.dtype):
279 sctype = _torch_dtypes[arg]
280 elif isinstance(arg, torch.Tensor):
281 sctype = _torch_dtypes[arg.dtype]
282 # a scalar type?
283 elif issubclass_(arg, generic):
284 sctype = arg
285 # a dtype already?
286 elif isinstance(arg, DType):
287 sctype = arg._scalar_type
288 # a has a right attribute?
289 elif hasattr(arg, "dtype"):
290 sctype = arg.dtype._scalar_type
291 else:
292 sctype = sctype_from_string(arg)
293 self._scalar_type = sctype
294
295 @property
296 def name(self):
297 return self._scalar_type.name
298
299 @property
300 def type(self):
301 return self._scalar_type
302
303 @property
304 def kind(self):
305 # https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
306 return _torch_dtypes[self.torch_dtype].name[0]
307
308 @property
309 def typecode(self):
310 return self._scalar_type.typecode
311
312 def __eq__(self, other):
313 if isinstance(other, DType):
314 return self._scalar_type == other._scalar_type
315 try:
316 other_instance = DType(other)
317 except TypeError:
318 return False
319 return self._scalar_type == other_instance._scalar_type
320
321 @property
322 def torch_dtype(self):
323 return self._scalar_type.torch_dtype
324
325 def __hash__(self):
326 return hash(self._scalar_type.name)
327
328 def __repr__(self):
329 return f'dtype("{self.name}")'
330
331 __str__ = __repr__
332
333 @property
334 def itemsize(self):
335 elem = self.type(1)
336 return elem.tensor.element_size()
337
338 def __getstate__(self):
339 return self._scalar_type
340
341 def __setstate__(self, value):
342 self._scalar_type = value
343
344
345typecodes = {
346 "All": "efdFDBbhil?",
347 "AllFloat": "efdFD",
348 "AllInteger": "Bbhil",
349 "Integer": "bhil",
350 "UnsignedInteger": "B",
351 "Float": "efd",
352 "Complex": "FD",
353}
354
355
356# ### Defaults and dtype discovery
357
358
359def set_default_dtype(fp_dtype="numpy", int_dtype="numpy"):
360 """Set the (global) defaults for fp, complex, and int dtypes.
361
362 The complex dtype is inferred from the float (fp) dtype. It has
363 a width at least twice the width of the float dtype,
364 i.e., it's complex128 for float64 and complex64 for float32.
365
366 Parameters
367 ----------
368 fp_dtype
369 Allowed values are "numpy", "pytorch" or dtype_like things which
370 can be converted into a DType instance.
371 Default is "numpy" (i.e. float64).
372 int_dtype
373 Allowed values are "numpy", "pytorch" or dtype_like things which
374 can be converted into a DType instance.
375 Default is "numpy" (i.e. int64).
376
377 Returns
378 -------
379 The old default dtype state: a namedtuple with attributes ``float_dtype``,
380 ``complex_dtypes`` and ``int_dtype``. These attributes store *pytorch*
381 dtypes.
382
383 Notes
384 ------------
385 This functions has a side effect: it sets the global state with the provided dtypes.
386
387 The complex dtype has bit width of at least twice the width of the float
388 dtype, i.e. it's complex128 for float64 and complex64 for float32.
389
390 """
391 if fp_dtype not in ["numpy", "pytorch"]:
392 fp_dtype = dtype(fp_dtype).torch_dtype
393 if int_dtype not in ["numpy", "pytorch"]:
394 int_dtype = dtype(int_dtype).torch_dtype
395
396 if fp_dtype == "numpy":
397 float_dtype = torch.float64
398 elif fp_dtype == "pytorch":
399 float_dtype = torch.float32
400 else:
401 float_dtype = fp_dtype
402
403 complex_dtype = {
404 torch.float64: torch.complex128,
405 torch.float32: torch.complex64,
406 torch.float16: torch.complex64,
407 }[float_dtype]
408
409 if int_dtype in ["numpy", "pytorch"]:
410 int_dtype = torch.int64
411
412 new_defaults = _dtypes_impl.DefaultDTypes(
413 float_dtype=float_dtype, complex_dtype=complex_dtype, int_dtype=int_dtype
414 )
415
416 # set the new global state and return the old state
417 old_defaults = _dtypes_impl.default_dtypes
418 _dtypes_impl._default_dtypes = new_defaults
419 return old_defaults
420
421
422def issubclass_(arg, klass):
423 try:
424 return issubclass(arg, klass)
425 except TypeError:
426 return False
427
428
429def issubdtype(arg1, arg2):
430 # cf https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numerictypes.py#L356-L420
431
432 # We also accept strings even if NumPy doesn't as dtypes are serialized as their
433 # string representation in dynamo's graph
434 def str_to_abstract(t):
435 if isinstance(t, str) and t in _abstract_dtypes:
436 return globals()[t]
437 return t
438
439 arg1 = str_to_abstract(arg1)
440 arg2 = str_to_abstract(arg2)
441
442 if not issubclass_(arg1, generic):
443 arg1 = dtype(arg1).type
444 if not issubclass_(arg2, generic):
445 arg2 = dtype(arg2).type
446 return issubclass(arg1, arg2)
447
448
449__all__ = ["dtype", "DType", "typecodes", "issubdtype", "set_default_dtype", "sctypes"]
450__all__ += list(_names.keys()) # noqa: PLE0605
451__all__ += list(_name_aliases.keys()) # noqa: PLE0605
452__all__ += _abstract_dtypes # noqa: PLE0605
453