chore: import upstream snapshot with attribution

python/tvm/topi/_te_tensor_overload.py

@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Register TOPI implementations for TE tensor overload hooks."""
+
+from tvm import te
+from tvm.te import _te_tensor_overload as _overload
+from tvm.tirx import expr as _expr
+
+from . import broadcast as _broadcast
+from . import math as _math
+
+
+def _is_integer(value):
+    if isinstance(value, te.Tensor | te.TensorSlice):
+        return value.dtype.matches_code(_expr.DataTypeCode.INT)
+    return _expr._dtype_is_int(value)
+
+
+def _binary(op, reflected=False, check_integer=False):
+    def implementation(lhs, rhs):
+        if not isinstance(lhs, te.Tensor) and not isinstance(rhs, te.Tensor):
+            return NotImplemented
+        if reflected:
+            lhs, rhs = rhs, lhs
+        if check_integer and _is_integer(lhs) and _is_integer(rhs):
+            raise _expr.div_ambiguity_error()
+        return op(lhs, rhs)
+
+    return implementation
+
+
+_overload.__add__ = _binary(_broadcast.add)
+_overload.__radd__ = _binary(_broadcast.add, reflected=True)
+_overload.__sub__ = _binary(_broadcast.subtract)
+_overload.__rsub__ = _binary(_broadcast.subtract, reflected=True)
+_overload.__mul__ = _binary(_broadcast.multiply)
+_overload.__rmul__ = _binary(_broadcast.multiply, reflected=True)
+_overload.__div__ = _binary(_broadcast.divide, check_integer=True)
+_overload.__rdiv__ = _binary(_broadcast.divide, reflected=True, check_integer=True)
+_overload.__truediv__ = _overload.__div__
+_overload.__rtruediv__ = _overload.__rdiv__
+
+
+def _astype(value, dtype, span=None):
+    if not isinstance(value, te.Tensor):
+        return NotImplemented
+    return _math.cast(value, dtype, span)
+
+
+_overload.astype = _astype