bbqddt2/Antigravity

1# mypy: ignore-errors
2
3"""Export torch work functions for binary ufuncs, rename/tweak to match numpy.
4This listing is further exported to public symbols in the `torch._numpy/_ufuncs.py` module.
5"""
6
7import torch
8from torch import (  # noqa: F401
9    add,
10    arctan2,
11    bitwise_and,
12    bitwise_left_shift as left_shift,
13    bitwise_or,
14    bitwise_right_shift as right_shift,
15    bitwise_xor,
16    copysign,
17    divide,
18    eq as equal,
19    float_power,
20    floor_divide,
21    fmax,
22    fmin,
23    fmod,
24    gcd,
25    greater,
26    greater_equal,
27    heaviside,
28    hypot,
29    lcm,
30    ldexp,
31    less,
32    less_equal,
33    logaddexp,
34    logaddexp2,
35    logical_and,
36    logical_or,
37    logical_xor,
38    maximum,
39    minimum,
40    multiply,
41    nextafter,
42    not_equal,
43    pow as power,
44    remainder,
45    remainder as mod,
46    subtract,
47    true_divide,
48)
49
50from . import _dtypes_impl, _util
51
52
53# work around torch limitations w.r.t. numpy
54def matmul(x, y):
55    # work around:
56    #  - RuntimeError: expected scalar type Int but found Double
57    #  - RuntimeError: "addmm_impl_cpu_" not implemented for 'Bool'
58    #  - RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
59    dtype = _dtypes_impl.result_type_impl(x, y)
60    is_bool = dtype == torch.bool
61    is_half = (x.dtype == torch.float16 or y.dtype == torch.float16) and (
62        x.is_cpu or y.is_cpu
63    )
64
65    work_dtype = dtype
66    if is_bool:
67        work_dtype = torch.uint8
68    if is_half:
69        work_dtype = torch.float32
70
71    x = _util.cast_if_needed(x, work_dtype)
72    y = _util.cast_if_needed(y, work_dtype)
73
74    result = torch.matmul(x, y)
75
76    if work_dtype != dtype:
77        result = result.to(dtype)
78
79    return result
80
81
82# a stub implementation of divmod, should be improved after
83# https://github.com/pytorch/pytorch/issues/90820 is fixed in pytorch
84def divmod(x, y):
85    return x // y, x % y
86