Add files using upload-large-folder tool

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/split_17_18.h

@@ -0,0 +1,38 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Split in default domain from version 17 to 18
+
+#pragma once
+
+#include <memory>
+
+#include "onnx/version_converter/adapters/adapter.h"
+#include "onnx/version_converter/adapters/transformers.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class Split_17_18 : public Adapter {
+ public:
+  explicit Split_17_18() : Adapter("Split", OpSetID(17), OpSetID(18)) {}
+
+  void adapt_split_17_18(const std::shared_ptr<Graph>&, Node* node) const {
+    const auto num_outputs = node->outputs().size();
+    node->i_(knum_outputs, num_outputs);
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    // if node does not have neither 'num_outputs' attribute nor 'split' input
+    if (!node->hasAttribute(knum_outputs) && node->inputs().size() != 2) {
+      adapt_split_17_18(graph, node);
+    }
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/sum_8_7.h

@@ -0,0 +1,41 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Sum in default domain from version 8 to 7
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class Sum_8_7 final : public Adapter {
+ public:
+  explicit Sum_8_7() : Adapter("Sum", OpSetID(8), OpSetID(7)) {}
+
+  void adapt_sum_8_7(const std::shared_ptr<Graph>&, Node* node) const {
+    // Throw an exception if any broadcasting occurs
+    const ArrayRef<Value*>& inputs = node->inputs();
+    // Determine if inputs are of different sizes
+    for (int i = 1; i < (int)inputs.size(); i++) {
+      std::vector<Dimension> A_sizes = inputs[i - 1]->sizes();
+      std::vector<Dimension> B_sizes = inputs[i]->sizes();
+      assert_numpy_multibroadcastable(A_sizes, B_sizes);
+    }
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_sum_8_7(graph, node);
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/topk_9_10.h

@@ -0,0 +1,45 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for TopK in default domain from version 9 to 10
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class TopK_9_10 final : public Adapter {
+ public:
+  explicit TopK_9_10() : Adapter("TopK", OpSetID(9), OpSetID(10)) {}
+
+  void adapt_topk_9_10(const std::shared_ptr<Graph>& graph, Node* node) const {
+    Tensor t;
+    t.elem_type() = TensorProto_DataType_INT64;
+    t.sizes() = std::vector<int64_t>{1};
+    auto& data = t.int64s();
+    data.emplace_back(node->i(kk));
+
+    Node* constant = graph->create(kConstant);
+    constant->insertBefore(node);
+    constant->t_(kvalue, t);
+    node->addInput(constant->output());
+
+    node->removeAttribute(kk);
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_topk_9_10(graph, node);
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/transformers.h

@@ -0,0 +1,85 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cinttypes>
+#include <string>
+#include <vector>
+
+#include "onnx/common/interned_strings.h"
+#include "onnx/version_converter/adapters/adapter.h"
+
+// Node transformers commonly used in version-adapters:
+
+// Capture context by copying values; the graph is unused by these transformers.
+
+#define NODE_TRANSFORMER(node) [=](const std::shared_ptr<Graph>&, Node* node)
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+inline NodeTransformerFunction RemoveAttribute(Symbol attr) {
+  return NODE_TRANSFORMER(node) {
+    if (node->hasAttribute(attr)) {
+      node->removeAttribute(attr);
+    }
+    return node;
+  };
+}
+
+inline NodeTransformerFunction RemoveAttribute(Symbol attr, int64_t value) {
+  return NODE_TRANSFORMER(node) {
+    if (node->hasAttribute(attr)) {
+      ONNX_ASSERTM(node->i(attr) == value, "Attribute %s must have value %" PRId64, attr.toString(), value)
+      node->removeAttribute(attr);
+    }
+    return node;
+  };
+}
+
+inline NodeTransformerFunction RemoveAttributeNotEq(Symbol attr, int64_t value) {
+  return NODE_TRANSFORMER(node) {
+    if (node->hasAttribute(attr)) {
+      ONNX_ASSERTM(node->i(attr) != value, "Attribute %s must not have value %" PRId64, attr.toString(), value)
+      node->removeAttribute(attr);
+    }
+    return node;
+  };
+}
+
+inline NodeTransformerFunction SetAttribute(Symbol attr, int64_t value) {
+  return NODE_TRANSFORMER(node) {
+    node->i_(attr, value);
+    return node;
+  };
+}
+
+inline NodeTransformerFunction SetAttribute(Symbol attr, const std::string& value) {
+  return NODE_TRANSFORMER(node) {
+    node->s_(attr, value);
+    return node;
+  };
+}
+
+inline NodeTransformerFunction SetAttribute(Symbol attr, const std::vector<int64_t>& value) {
+  return NODE_TRANSFORMER(node) {
+    node->is_(attr, std::vector<int64_t>(value));
+    return node;
+  };
+}
+
+inline NodeTransformerFunction SetAttributeIfAbsent(Symbol attr, int64_t value) {
+  return NODE_TRANSFORMER(node) {
+    if (!node->hasAttribute(attr)) {
+      node->i_(attr, value);
+    }
+    return node;
+  };
+}
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/type_restriction.h

@@ -0,0 +1,61 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Add in default domain from version 6 to 5
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class TypeRestriction : public Adapter {
+ public:
+  explicit TypeRestriction(
+      const std::string& op_name,
+      const OpSetID& initial,
+      const OpSetID& target,
+      const std::vector<TensorProto_DataType>& unallowed_types)
+      : Adapter(op_name, initial, target), unallowed_types_(unallowed_types) {}
+
+  void adapt_type_restriction(const std::shared_ptr<Graph>&, const Node* node) const {
+    // Since consumed_inputs is optional, no need to add it (as in batchnorm)
+    // Iterate over all inputs and outputs
+    for (const Value* input : node->inputs()) {
+      isUnallowed(input);
+    }
+    for (const Value* output : node->outputs()) {
+      isUnallowed(output);
+    }
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_type_restriction(graph, node);
+    return node;
+  }
+
+ private:
+  std::vector<TensorProto_DataType> unallowed_types_;
+
+  void isUnallowed(const Value* val) const {
+    ONNX_ASSERTM(
+        std::find(std::begin(unallowed_types_), std::end(unallowed_types_), val->elemType()) ==
+            std::end(unallowed_types_),
+        "DataType (%d) of Input or Output"
+        " of operator '%s' is unallowed for Opset Version %d.",
+        val->elemType(),
+        name().c_str(),
+        target_version().version())
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/upsample_6_7.h

@@ -0,0 +1,49 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Upsample in default domain from version 6 to 7
+
+#pragma once
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+struct Upsample_6_7 final : public Adapter {
+  explicit Upsample_6_7() : Adapter("Upsample", OpSetID(6), OpSetID(7)) {}
+
+  void adapt_upsample_6_7(const std::shared_ptr<Graph>&, Node* node) const {
+    Symbol width_scale_symbol = Symbol("width_scale");
+    Symbol height_scale_symbol = Symbol("height_scale");
+    ONNX_ASSERTM(
+        node->hasAttribute(width_scale_symbol) && node->hasAttribute(height_scale_symbol),
+        "Upsample in opset 1 needs to have width_scale and height_scale attributes")
+
+    auto width_scale = node->f(width_scale_symbol);
+    auto height_scale = node->f(height_scale_symbol);
+
+    auto input_shape = node->inputs()[0]->sizes();
+    ONNX_ASSERTM(input_shape.size() == 4, "Upsample in opset 1 supports only 4D input tensor")
+    std::vector<double> scales = {1.0, 1.0, height_scale, width_scale};
+
+    node->fs_(kscales, std::move(scales));
+    node->removeAttribute(width_scale_symbol);
+    node->removeAttribute(height_scale_symbol);
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_upsample_6_7(graph, node);
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/upsample_8_9.h

@@ -0,0 +1,50 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Upsample in default domain from version 8 to 9
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+struct Upsample_8_9 final : public Adapter {
+  explicit Upsample_8_9() : Adapter("Upsample", OpSetID(8), OpSetID(9)) {}
+
+  void adapt_upsample_8_9(const std::shared_ptr<Graph>& graph, Node* node) const {
+    Symbol input_dirs = Symbol("scales");
+    int dim = (int)(node->fs(kscales).size());
+    Tensor t;
+    t.elem_type() = TensorProto_DataType_FLOAT;
+    t.sizes() = std::vector<int64_t>{dim};
+    auto& data = t.floats();
+
+    if (node->hasAttribute(input_dirs)) {
+      for (double scale : node->fs(kscales)) {
+        data.emplace_back((float)scale);
+      }
+
+      Node* constant = graph->create(kConstant);
+      constant->insertBefore(node);
+      constant->t_(kvalue, t);
+      node->addInput(constant->output());
+      node->removeAttribute(kscales);
+    }
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_upsample_8_9(graph, node);
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/upsample_9_10.h

@@ -0,0 +1,45 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Upsample in default domain from version 9 to 10
+
+#pragma once
+
+#include <memory>
+#include <string>
+
+#include "onnx/common/ir.h"
+#include "onnx/version_converter/adapters/adapter.h"
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class Upsample_9_10 final : public Adapter {
+ public:
+  explicit Upsample_9_10() : Adapter("Upsample", OpSetID(9), OpSetID(10)) {}
+
+  Node* adapt_upsample_9_10(const std::shared_ptr<Graph>& graph, Node* node) const {
+    std::string mode = node->hasAttribute(kmode) ? node->s(kmode) : "nearest";
+
+    // Replace the node with an equivalent Resize node
+    Node* resize = graph->create(kResize);
+    resize->s_(kmode, mode);
+    resize->addInput(node->inputs()[0]);
+    resize->addInput(node->inputs()[1]);
+    node->replaceAllUsesWith(resize);
+
+    resize->insertBefore(node);
+    node->destroy();
+
+    return resize;
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    return adapt_upsample_9_10(graph, node);
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnx/version_converter/adapters/upsample_9_8.h

@@ -0,0 +1,79 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Adapter for Upsample in default domain from version 9 to 8
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "onnx/defs/tensor_proto_util.h"
+#include "onnx/defs/tensor_util.h"
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+struct Upsample_9_8 final : public Adapter {
+  explicit Upsample_9_8() : Adapter("Upsample", OpSetID(9), OpSetID(8)) {}
+
+  void adapt_upsample_9_8(const std::shared_ptr<Graph>& graph, Node* node) const {
+    const ArrayRef<Value*>& inputs = node->inputs();
+    const std::vector<Tensor>& initializers = graph->initializers();
+
+    ONNX_ASSERTM(inputs.size() == 2, "Upsample in opset 9 needs to have 2 inputs.")
+    std::string scale_input_name = node->inputs()[1]->uniqueName();
+
+    for (const auto& initializer : initializers) {
+      if (initializer.name() == inputs[1]->uniqueName()) {
+        std::vector<float> value = ParseData<float>(&initializer);
+        std::vector<double> d_values;
+        d_values.reserve(value.size());
+        for (float j : value) {
+          d_values.push_back(static_cast<double>(j));
+        }
+        node->fs_(kscales, std::move(d_values));
+
+        node->removeInput(1);
+        graph->eraseInitializer(initializer.name());
+        for (size_t j = 0; j < graph->inputs().size(); j++) {
+          if (graph->inputs()[j]->uniqueName() == scale_input_name) {
+            graph->eraseInput(j);
+            break;
+          }
+        }
+        return;
+      }
+    }
+
+    for (Node* op : graph->nodes()) {
+      if (op->kind() == kConstant && op->outputs()[0]->uniqueName() == scale_input_name) {
+        std::vector<float> value = ParseData<float>(&op->t(kvalue));
+        std::vector<double> d_values;
+        d_values.reserve(value.size());
+        for (float j : value) {
+          d_values.push_back(static_cast<double>(j));
+        }
+        node->fs_(kscales, std::move(d_values));
+        node->removeInput(1);
+        op->destroy();
+        return;
+      }
+    }
+
+    ONNX_ASSERTM(false, "Unsuppported conversion due to unavailable input: scale")
+  }
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    adapt_upsample_9_8(graph, node);
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

pythonProject/.venv/Lib/site-packages/onnxscript/function_libs/torch_lib/_constants.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Shared constants for the library."""
+
+DOMAIN = "pkg.onnxscript.torch_lib"

pythonProject/.venv/Lib/site-packages/onnxscript/function_libs/torch_lib/_flags.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Experimental flags.
+
+NOTE: These flags are experimental only. Any flag here can be removed at any
+time without notice.
+"""
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+
+def _load_boolean_flag(
+    name: str,
+    *,
+    this_will: str,
+    deprecated: bool = False,
+    default: bool = False,
+) -> bool:
+    """Load a boolean flag from environment variable.
+
+    Args:
+        name: The name of the environment variable.
+        this_will: A string that describes what this flag will do.
+        deprecated: Whether this flag is deprecated.
+        default: The default value if envvar not defined.
+    """
+    undefined = os.getenv(name) is None
+    state = os.getenv(name) == "1"
+    if state:
+        if deprecated:
+            logger.error(
+                "Experimental flag %s is deprecated. Please remove it from your environment.",
+                name,
+            )
+        else:
+            logger.warning("Experimental flag %s is enabled. This will %s.", name, this_will)
+    if undefined:
+        state = default
+    return state
+
+
+EXPERIMENTAL_INITIALIZERS_AS_INPUTS: bool = _load_boolean_flag(
+    "TORCHLIB_EXPERIMENTAL_INITIALIZERS_AS_INPUTS",
+    this_will="make initializers as inputs to the model graph",
+)
+EXPERIMENTAL_PREFER_TRACING: bool = _load_boolean_flag(
+    "TORCHLIB_EXPERIMENTAL_PREFER_TRACING",
+    this_will="trace all traceable functions to fold if branches and collapse constant expressions",
+    default=True,
+)

pythonProject/.venv/Lib/site-packages/onnxscript/function_libs/torch_lib/registration.py

@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Registry for aten functions."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Callable, Generator, Optional
+
+import onnxscript
+from onnxscript.function_libs.torch_lib import _constants
+
+# Regex that will match "<namespace>::<op_name>[.<overload>]"
+_QUALIFIED_OPERATOR_NAME_REGEX = re.compile(
+    r"^(?P<namespace>[a-zA-Z0-9_]+)::(?P<name>[a-zA-Z0-9_]+)(?P<overload>\.[a-zA-Z0-9._]+)?$"
+)
+
+
+class OverloadedFunction:
+    """Overloaded function.
+
+    Attributes:
+        name: Name of the op. E.g. "aten::add".
+        overloads: Overloads function.
+        privates: Private functions not exposed to users.
+        complex: Support complex functions.
+    """
+
+    def __init__(self, name: str):
+        self.name = name
+        self.overloads: list[Any] = []
+        self.privates: list[Any] = []
+        self.complex: list[Any] = []
+
+
+class Registry:
+    """Registry for aten functions."""
+
+    def __init__(self):
+        self._registry: dict[str, OverloadedFunction] = {}
+
+    def register(
+        self, func: Any, name: str, *, private: bool = False, complex: bool = False
+    ) -> None:
+        """Register a function."""
+
+        if private:
+            self._registry.setdefault(name, OverloadedFunction(name)).privates.append(func)
+        elif complex:
+            self._registry.setdefault(name, OverloadedFunction(name)).complex.append(func)
+        else:
+            self._registry.setdefault(name, OverloadedFunction(name)).overloads.append(func)
+
+    def __getitem__(self, name):
+        return self._registry[name]
+
+    def __contains__(self, name):
+        return name in self._registry
+
+    def __iter__(self):
+        return iter(self._registry)
+
+    def __repr__(self):
+        return repr(self._registry)
+
+    def items(self) -> Generator[tuple[str, OverloadedFunction], None, None]:
+        yield from self._registry.items()
+
+    def values(self) -> Generator[OverloadedFunction, None, None]:
+        yield from self._registry.values()
+
+
+# Default registry
+default_registry = Registry()
+
+
+def _check_and_normalize_names(name: str | tuple[str, ...]) -> tuple[str, ...]:
+    names: tuple[str, ...]
+
+    if isinstance(name, str):
+        names = (name,)
+    else:
+        names = name
+    if not isinstance(names, tuple):
+        raise TypeError(f"Name must be a string or a tuple of strings, got {name}")
+    for name_ in names:
+        if name_.endswith(".default") or not _QUALIFIED_OPERATOR_NAME_REGEX.fullmatch(name_):
+            raise ValueError(
+                f"Invalid name '{name_}'. Must be in the form 'namespace::name' for default overloads "
+                "or 'namespace::name.overload' for other overloads."
+            )
+
+    return names
+
+
+def torch_op(
+    name: str | tuple[str, ...],
+    *,
+    registry: Optional[Registry] = None,
+    trace_only: bool = False,
+    private: bool = False,
+    complex: bool = False,
+) -> Callable[[Callable], onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction]:
+    """Register a torch op.
+
+    Args:
+        name: Qualified ATen name of the function. E.g. "aten::relu", "aten::add.Tensor".
+            Or a tuple of names e.g. ("aten::add.Scalar", "aten::add.Tensor").
+            Default overloads should be specified by omitting the overload part,
+            i.e. "aten::relu" instead of "aten::relu.default".
+        registry: Registry to register the function to. If None, the default registry is used.
+        trace_only: Whether the function should only be traced and not compiled.
+        private: Whether the function is private (not directly exposed). It should
+            be true for all functions with names starting with "_".
+        complex: Whether the function expects complex-valued inputs.
+    """
+    if registry is None:
+        registry = default_registry
+
+    def wrapper(
+        func: Callable,
+    ) -> onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction:
+        # Compile the function
+        custom_opset = onnxscript.values.Opset(domain=_constants.DOMAIN, version=1)
+
+        processed_func: onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction
+        if trace_only:
+            processed_func = onnxscript.values.TracedOnnxFunction(custom_opset, func)
+        else:
+            processed_func = onnxscript.script(opset=custom_opset)(func)
+
+        assert registry is not None
+        for name_ in _check_and_normalize_names(name):
+            registry.register(processed_func, name_, private=private, complex=complex)
+        return processed_func
+
+    return wrapper

pythonProject/.venv/Lib/site-packages/onnxscript/function_libs/torch_lib/tensor_typing.py

@@ -0,0 +1,74 @@
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""Typings for function definitions."""
+
+from __future__ import annotations
+
+from typing import TypeVar, Union
+
+from onnxscript import (
+    BFLOAT16,
+    BOOL,
+    COMPLEX64,
+    COMPLEX128,
+    DOUBLE,
+    FLOAT,
+    FLOAT16,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    STRING,
+    UINT8,
+)
+
+# NOTE: We do not care about unsigned types beyond UINT8 because PyTorch does not us them.
+# More detail can be found: https://pytorch.org/docs/stable/tensors.html
+
+_TensorType = Union[
+    BFLOAT16,
+    BOOL,
+    COMPLEX64,
+    COMPLEX128,
+    DOUBLE,
+    FLOAT,
+    FLOAT16,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    UINT8,
+]
+_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]
+IntType = Union[INT8, INT16, INT32, INT64]
+RealType = Union[
+    BFLOAT16,
+    FLOAT16,
+    FLOAT,
+    DOUBLE,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+]
+
+TTensor = TypeVar("TTensor", bound=_TensorType)
+# Duplicate TTensor for inputs/outputs that accept the same set of types as TTensor
+# but do not constrain the type to be the same as the other inputs/outputs
+TTensor2 = TypeVar("TTensor2", bound=_TensorType)
+TTensorOrString = TypeVar("TTensorOrString", bound=Union[_TensorType, STRING])
+TFloat = TypeVar("TFloat", bound=_FloatType)
+TFloatOrUInt8 = TypeVar("TFloatOrUInt8", bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8])
+TInt = TypeVar("TInt", bound=IntType)
+TReal = TypeVar("TReal", bound=RealType)
+TRealUnlessInt16OrInt8 = TypeVar(
+    "TRealUnlessInt16OrInt8", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64]
+)
+TRealUnlessFloat16OrInt8 = TypeVar(
+    "TRealUnlessFloat16OrInt8", bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64]
+)
+TRealOrUInt8 = TypeVar("TRealOrUInt8", bound=Union[RealType, UINT8])
+TFloatHighPrecision = TypeVar("TFloatHighPrecision", bound=Union[FLOAT, DOUBLE])