Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/_library/__init__.py

@@ -0,0 +1,6 @@
+import torch._library.autograd
+import torch._library.fake_impl
+import torch._library.simple_registry
+import torch._library.utils
+from torch._library.fake_class_registry import register_fake_class
+from torch._library.triton import capture_triton, triton_op, wrap_triton

kronos-venv/lib/python3.12/site-packages/torch/_library/_out_variant.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import logging
+
+import torch
+
+
+log = logging.getLogger(__name__)
+
+
+# Manual registry for ops whose out variant is not discoverable via
+# to_out_variant() (e.g. flat _out naming instead of .out overload).
+_manual_out_variant_registry: dict[torch._ops.OpOverload, torch._ops.OpOverload] = {}
+
+
+def register_out_variant(
+    functional_op: torch._ops.OpOverload,
+    out_op: torch._ops.OpOverload,
+) -> None:
+    """Register a functional op -> out variant mapping."""
+    _manual_out_variant_registry[functional_op] = out_op
+
+
+def lookup_manual_out_variant(
+    op: torch._ops.OpOverload,
+) -> torch._ops.OpOverload | None:
+    """Return the manually registered out variant for op, or None."""
+    return _manual_out_variant_registry.get(op)
+
+
+def _is_functional(schema: torch._C.FunctionSchema) -> bool:
+    """
+    A schema is functional if no argument is written to and the name doesn't
+    end with '_'.
+    """
+    op_name = schema.name.split("::")[-1]
+    if op_name.endswith("_"):
+        return False
+    return not any(arg.is_write for arg in schema.arguments)
+
+
+def _is_mutable_arg(arg: torch._C.Argument) -> bool:
+    return arg.alias_info is not None and arg.alias_info.is_write
+
+
+def _signatures_match(
+    schema_a: torch._C.FunctionSchema,
+    schema_b: torch._C.FunctionSchema,
+) -> bool:
+    """Compare two schemas by their non-mutable arguments (name, type, default value)."""
+    non_mutable_args_a = [arg for arg in schema_a.arguments if not _is_mutable_arg(arg)]
+    non_mutable_args_b = [arg for arg in schema_b.arguments if not _is_mutable_arg(arg)]
+    if len(non_mutable_args_a) != len(non_mutable_args_b):
+        return False
+    for a, b in zip(non_mutable_args_a, non_mutable_args_b):
+        if a.name != b.name:
+            return False
+        if str(a.type) != str(b.type):
+            return False
+        if a.default_value != b.default_value:
+            return False
+    return True
+
+
+def _has_valid_out_variant_returns(
+    schema: torch._C.FunctionSchema,
+    mutable_args: list[torch._C.Argument],
+) -> bool:
+    """Out variant must return either nothing or the mutable args themselves."""
+    if len(schema.returns) == 0:
+        return True
+
+    if len(schema.returns) != len(mutable_args):
+        return False
+
+    # Each return must alias exactly one mutable arg, in order
+    for ret, arg in zip(schema.returns, mutable_args):
+        if ret.alias_info is None or arg.alias_info is None:
+            return False
+        if ret.alias_info.before_set != arg.alias_info.before_set:
+            return False
+    return True
+
+
+def get_out_arg_names(out_op: torch._ops.OpOverload) -> list[str]:
+    """Get the names of out arguments for an out variant op."""
+    schema = out_op._schema
+    return [arg.name for arg in schema.arguments if _is_mutable_arg(arg)]
+
+
+def to_out_variant(op: torch._ops.OpOverload) -> torch._ops.OpOverload | None:
+    """
+    Given a functional operator overload, return its corresponding out variant.
+    """
+    schema = op._schema
+
+    if not _is_functional(schema):
+        raise RuntimeError(
+            f"Failed to find out variant for op '{op}' as its schema is not functional. \n"
+            f"  {schema}"
+        )
+
+    # Get the op packet to access all overloads
+    namespace = op.namespace
+    op_name = schema.name.split("::")[1]
+    torch_packet = getattr(getattr(torch.ops, namespace), op_name)
+
+    # Search through all overloads for matching out variant
+    for overload_name in torch_packet.overloads():
+        candidate = getattr(torch_packet, overload_name)
+
+        # pyrefly: ignore [missing-attribute]
+        if torch.Tag.out_variant not in candidate.tags:
+            continue
+
+        candidate_schema = candidate._schema
+
+        if not _signatures_match(schema, candidate_schema):
+            continue
+
+        # We assume that all mutable args are used for out
+        mutable_args = [
+            arg for arg in candidate_schema.arguments if _is_mutable_arg(arg)
+        ]
+        if len(mutable_args) != len(schema.returns):
+            continue
+
+        if not _has_valid_out_variant_returns(candidate_schema, mutable_args):
+            raise RuntimeError(
+                f"Out variant {candidate} has invalid returns. "
+                f"Expected either no returns or returns that alias the mutable args, "
+                f"got: {candidate_schema}"
+            )
+
+        return candidate
+
+    return None
+
+
+def check_out_variant(
+    functional_op: torch._ops.OpOverload, expected_out_op: torch._ops.OpOverload
+) -> None:
+    """
+    Checks that to_out_variant returns the expected out variant for a functional op.
+    Raises AssertionError if the out variant is not valid.
+    """
+    out_op = to_out_variant(functional_op)
+    if out_op is None:
+        tagged_info = _get_out_variants_info(functional_op)
+        raise AssertionError(
+            f"We did not find an out variant for {functional_op}. Some common mistakes include:\n"
+            "  1. The out variant is missing the torch.Tag.out_variant tag.\n"
+            "  2. The out variant is not an overload of the original op (e.g., 'op.out' or 'op.overload_out') \n"
+            "  3. The out variant's input arguments does not match the functional op's signature (excluding the mutable args).\n"
+            "  4. The original operator is not functional.\n"
+            f"Overloads tagged with out_variant:\n"
+            f"{tagged_info or '  (none)'}"
+        )
+    if out_op != expected_out_op:
+        raise AssertionError(
+            f"to_out_variant({functional_op}) returned {out_op}, "
+            f"but expected {expected_out_op}. "
+            f"The out variant name does not match the functional op."
+        )
+
+
+def _get_out_variants_info(functional_op) -> str:
+    """Collect information about overloads tagged with out_variant for debugging."""
+    namespace = functional_op.namespace
+    op_name = functional_op._schema.name.split("::")[1]
+    torch_packet = getattr(getattr(torch.ops, namespace), op_name)
+
+    overloads_info: list[str] = []
+    for overload_name in torch_packet.overloads():
+        candidate = getattr(torch_packet, overload_name)
+        # pyrefly: ignore [missing-attribute]
+        if torch.Tag.out_variant in candidate.tags:
+            overloads_info.append(f"  - {overload_name}: {candidate._schema}")
+
+    return "\n".join(overloads_info)

kronos-venv/lib/python3.12/site-packages/torch/_library/autograd.py

@@ -0,0 +1,240 @@
+# mypy: allow-untyped-defs
+import dataclasses
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from torch import _C, _ops, autograd, Tensor
+from torch.utils import _pytree
+
+from . import utils
+
+
+class InfoProtocol(Protocol):
+    _backward_fn: Callable | None
+    _setup_context_fn: Callable | None
+
+
+@dataclasses.dataclass
+class Info:
+    _backward_fn: Callable | None
+    _setup_context_fn: Callable | None
+
+
+def make_autograd_impl(op: _ops.OpOverload, info: InfoProtocol) -> Callable:
+    name: str = f"GeneratedBackwardFor_{op._namespace}_{op._opname}_{op._overloadname}"
+
+    has_kwarg_only_args = utils.has_kwarg_only_args(op._schema)
+
+    @dataclass
+    class Metadata:
+        keyset: _C.DispatchKeySet
+        keyword_only_args: dict[str, Any]
+
+    def forward_no_grad(*args):
+        metadata = args[-1]
+        args = args[:-1]
+
+        with _C._AutoDispatchBelowAutograd():
+            keyset = metadata.keyset
+            kwargs = metadata.keyword_only_args
+            result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
+            return result
+
+    def forward(ctx, *args):
+        metadata = args[-1]
+        args = args[:-1]
+
+        with _C._AutoDispatchBelowAutograd():
+            keyset = metadata.keyset
+            kwargs = metadata.keyword_only_args
+            result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
+            if info._setup_context_fn:
+                # The Dispatcher will remove args that are equal to their default
+                # values from (args, kwargs). We're going to add it back so that
+                # the user can access them.
+                #
+                # This is OK to do: The Dispatcher removed the args for serialization
+                # FC/BC reasons (that is, a graph will not store args that are equal
+                # to their default values), but that doesn't matter here. If the user
+                # adds a new default arg, then they must update
+                # their setup_context (along with the rest of their operator
+                # registrations)
+                args, kwargs = utils.fill_defaults(op._schema, args, kwargs)
+
+                if has_kwarg_only_args:
+                    info._setup_context_fn(
+                        ctx=ctx, inputs=args, keyword_only_inputs=kwargs, output=result
+                    )
+                else:
+                    info._setup_context_fn(ctx=ctx, inputs=args, output=result)
+            return result
+
+    def backward(ctx, *grads):
+        if info._backward_fn:
+            try:
+                prev_needs_input_grad = ctx.needs_input_grad
+                ctx.needs_input_grad = ctx.needs_input_grad[:-1]
+                result = info._backward_fn(ctx, *grads)
+            finally:
+                ctx.needs_input_grad = prev_needs_input_grad
+            if isinstance(result, tuple):
+                return (*result, None)
+            return result, None
+        raise RuntimeError(
+            f"Trying to backward through {op} but no autograd "
+            f"formula was registered. "
+            f"Please use register_autograd to add one."
+        )
+
+    Generated = type(
+        name,
+        (autograd.Function,),
+        {
+            "forward": staticmethod(forward),
+            "backward": staticmethod(backward),
+        },
+    )
+
+    schema = op._schema
+    if any(
+        utils.is_tensorlist_like_type(a.type)
+        for a in (*schema.arguments, *schema.returns)
+    ):
+        Generated = supports_tensorlist(Generated)
+
+    # The dispatcher passes any keyword-only-args as kwargs and the
+    # rest of the args (even if specified as kwargs) as args.
+    def autograd_impl(keyset, *args, **keyword_only_args):
+        if _C.is_grad_enabled() and _C._any_requires_grad(*args):
+            result = Generated.apply(*args, Metadata(keyset, keyword_only_args))  # type: ignore[attr-defined]
+        else:
+            result = forward_no_grad(*args, Metadata(keyset, keyword_only_args))
+        return result
+
+    return autograd_impl
+
+
+def supports_tensorlist(cls: Any) -> Any:
+    """Allows a given autograd.Function class to support List[Tensor] inputs/outputs.
+
+    Regular autograd.Function has a constraint that it only directly supports autograd for
+    Tensors. Applying @supports_tensorlist enables an autograd.Function to support
+    autograd for List[Tensor] inputs and outputs.
+    """
+    orig_forward = cls.forward
+    orig_backward = cls.backward
+    orig_apply = cls.apply
+
+    @dataclass
+    class Metadata:
+        input_spec: _pytree.TreeSpec
+        output_spec: _pytree.TreeSpec | None = None
+        result_is_tuple: bool | None = None
+
+    def new_forward(ctx, *args):
+        metadata = args[-1]
+        args = args[:-1]
+        if not isinstance(metadata, Metadata):
+            raise NotImplementedError(
+                "NYI: calling supports_tensorlist autograd.Function.forward directly. "
+                "You should probably be calling .apply instead. "
+                "Please file an issue if not."
+            )
+        args = _pytree.tree_unflatten(list(args), metadata.input_spec)
+        result = orig_forward(ctx, *args)
+        metadata.result_is_tuple = isinstance(result, tuple)
+        if not metadata.result_is_tuple:
+            result = (result,)
+        flat_result, output_spec = _pytree.tree_flatten(result, not_list_of_tensor)
+        metadata.output_spec = output_spec
+
+        if hasattr(ctx, "_pt_metadata"):
+            raise RuntimeError(
+                "Please don't set ctx._pt_metadata; PyTorch uses it to store info"
+            )
+        ctx._pt_metadata = metadata
+
+        return tuple(flat_result)
+
+    def new_backward(ctx, *grads):
+        if not hasattr(ctx, "_pt_metadata"):
+            raise NotImplementedError(
+                "NYI: calling supports_tensorlist autograd.Function.backward directly. "
+                "This will automatically get called by PyTorch autograd. "
+                "Please file an issue if you need this."
+            )
+
+        metadata = ctx._pt_metadata
+        grads = _pytree.tree_unflatten(list(grads), metadata.output_spec)
+
+        # If the user's input is ([x, y, z], w),
+        # then needs_input_grad is (bool, bool, bool, bool, bool).
+        # We need to
+        # 1. get rid of the additional bool (which comes from the extra
+        # `metadata input`)
+        # 2. _pytree.tree_unflatten to get the right structure.
+        prev_needs_input_grad = ctx.needs_input_grad
+        try:
+            ctx.needs_input_grad = _pytree.tree_unflatten(
+                list(ctx.needs_input_grad[:-1]), metadata.input_spec
+            )
+            grad_inputs = orig_backward(ctx, *grads)
+        finally:
+            ctx.needs_input_grad = prev_needs_input_grad
+
+        if not isinstance(grad_inputs, tuple):
+            grad_inputs = (grad_inputs,)
+        # Assume that any Nones in the backward are Tensors.
+        # If the forward has an arg that is [1, 2, 3], the backward should
+        # return None as the grad.
+        # If the forward has an arg that is [tensor, tensor], the backward
+        # may return [None, None], [grad, None], [None, grad], or [grad, grad].
+        flat_grad_inputs, grad_inputs_spec = _pytree.tree_flatten(
+            grad_inputs, not_list_of_optional_tensor
+        )
+        if grad_inputs_spec != metadata.input_spec:
+            raise RuntimeError(
+                f"Expected the return from backward to be of the same structure "
+                f"as the inputs. Got: {grad_inputs_spec} (return from backward), "
+                f"{metadata.input_spec} (inputs)"
+            )
+        return tuple(flat_grad_inputs + [None])
+
+    def new_apply(*args):
+        flat_args, input_spec = _pytree.tree_flatten(args, is_leaf=not_list_of_tensor)
+        metadata = Metadata(input_spec)
+        result = orig_apply(*flat_args, metadata)  # type: ignore[misc]
+        if metadata.output_spec is None:
+            raise AssertionError("metadata.output_spec must not be None")
+        result = _pytree.tree_unflatten(list(result), metadata.output_spec)
+        if not metadata.result_is_tuple:
+            if not isinstance(result, tuple):
+                raise AssertionError(f"result must be tuple, got {type(result)}")
+            if len(result) != 1:
+                raise AssertionError(
+                    f"result tuple must have length 1, got {len(result)}"
+                )
+            return result[0]
+        return result
+
+    cls.forward = new_forward
+    cls.backward = new_backward
+    cls.apply = new_apply
+    return cls
+
+
+def not_list_of_tensor(tree):
+    if isinstance(tree, tuple):
+        return False
+    if isinstance(tree, list):
+        return any(not isinstance(l, Tensor) for l in tree)
+    return True
+
+
+def not_list_of_optional_tensor(tree):
+    if isinstance(tree, tuple):
+        return False
+    if isinstance(tree, list):
+        return any(l is not None and not isinstance(l, Tensor) for l in tree)
+    return True

kronos-venv/lib/python3.12/site-packages/torch/_library/custom_ops.py

@@ -0,0 +1,987 @@
+# mypy: allow-untyped-defs
+import collections
+import inspect
+import logging
+import warnings
+import weakref
+from collections.abc import Callable, Iterable, Sequence
+from contextlib import contextmanager
+from typing import Any, overload, Union
+
+import torch
+from torch import _C, _ops, Tensor
+from torch.types import _dtype
+from torch.utils._exposed_in import exposed_in
+
+from . import autograd, utils
+from .effects import EffectType
+
+
+device_types_t = str | Sequence[str] | None
+log = logging.getLogger(__name__)
+
+
+@overload
+def custom_op(
+    name: str,
+    fn: None = None,
+    /,
+    *,
+    mutates_args: str | Iterable[str],
+    device_types: device_types_t = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
+) -> Callable[[Callable[..., object]], "CustomOpDef"]: ...
+
+
+@overload
+def custom_op(
+    name: str,
+    fn: Callable[..., object],
+    /,
+    *,
+    mutates_args: str | Iterable[str],
+    device_types: device_types_t = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
+) -> "CustomOpDef": ...
+
+
+@exposed_in("torch.library")
+def custom_op(
+    name: str,
+    fn: Callable | None = None,
+    /,
+    *,
+    mutates_args: str | Iterable[str],
+    device_types: device_types_t = None,
+    schema: str | None = None,
+    tags: Sequence[_C.Tag] | None = None,
+) -> Union[Callable[[Callable[..., object]], "CustomOpDef"], "CustomOpDef"]:
+    """Wraps a function into custom operator.
+
+    Reasons why you may want to create a custom op include:
+    - Wrapping a third-party library or custom kernel to work with PyTorch
+    subsystems like Autograd.
+    - Preventing torch.compile/export/FX tracing from peeking inside your function.
+
+    This API is used as a decorator around a function (please see examples).
+    The provided function must have type hints; these are needed to interface
+    with PyTorch's various subsystems.
+
+    Args:
+        name (str): A name for the custom op that looks like "{namespace}::{name}",
+            e.g. "mylib::my_linear". The name is used as the op's stable identifier
+            in PyTorch subsystems (e.g. torch.export, FX graphs).
+            To avoid name collisions, please use your project name as the namespace;
+            e.g. all custom ops in pytorch/fbgemm use "fbgemm" as the namespace.
+        mutates_args (Iterable[str] or "unknown"): The names of args that the function mutates.
+            This MUST be accurate, otherwise, the behavior is undefined. If "unknown",
+            it pessimistically assumes that all inputs to the operator are being mutated.
+        device_types (str | None | Sequence[str]): The device type(s) the function
+            is valid for. If no device type is provided, then the function
+            is used as the default implementation for all device types.
+            Examples: "cpu", "cuda".
+            When registering a device-specific implementation for an operator that accepts no Tensors,
+            we require the operator to have a "device: torch.device argument".
+        schema (str | None): A schema string for the operator. If None
+            (recommended) we'll infer a schema for the operator from its type
+            annotations. We recommend letting us infer a schema unless you
+            have a specific reason not to.
+            Example: "(Tensor x, int y) -> (Tensor, Tensor)".
+
+    The following types are supported for the wrapped function's input parameters:
+
+        - Scalars: ``int``, ``float``, ``bool``, ``str``, ``torch.types.Number``
+        - Tensors: ``torch.Tensor``
+        - Enums/devices: ``torch.dtype``, ``torch.device``
+        - Flat list of the same type: ``list[torch.Tensor]``,
+          ``list[int]``, ``list[float]``, ``list[bool]``,
+          ``list[torch.types.Number]``
+        - Optionals: ``Optional`` of any of the above scalar/tensor types
+        - Types registered via :func:`torch.library.register_opaque_type`
+
+    The following types are supported for the return value:
+
+        ``torch.Tensor``, ``list[torch.Tensor]``, ``int``, ``float``,
+        ``bool``, ``torch.types.Number``.
+
+    .. note::
+        We recommend not passing in a ``schema`` arg and instead letting us infer
+        it from the type annotations. It is error-prone to write your own schema.
+        You may wish to provide your own schema if our interpretation of
+        the type annotation is not what you want.
+        For more info on how to write a schema string, see
+        `here <https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#func>`_
+
+    Examples::
+        >>> import torch
+        >>> from torch import Tensor
+        >>> from torch.library import custom_op
+        >>> import numpy as np
+        >>>
+        >>> @custom_op("mylib::numpy_sin", mutates_args=())
+        >>> def numpy_sin(x: Tensor) -> Tensor:
+        >>>     x_np = x.cpu().numpy()
+        >>>     y_np = np.sin(x_np)
+        >>>     return torch.from_numpy(y_np).to(device=x.device)
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = numpy_sin(x)
+        >>> assert torch.allclose(y, x.sin())
+        >>>
+        >>> # Example of a custom op that only works for one device type.
+        >>> @custom_op("mylib::numpy_sin_cpu", mutates_args=(), device_types="cpu")
+        >>> def numpy_sin_cpu(x: Tensor) -> Tensor:
+        >>>     x_np = x.numpy()
+        >>>     y_np = np.sin(x_np)
+        >>>     return torch.from_numpy(y_np)
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = numpy_sin_cpu(x)
+        >>> assert torch.allclose(y, x.sin())
+        >>>
+        >>> # Example of a custom op that mutates an input
+        >>> @custom_op("mylib::numpy_sin_inplace", mutates_args={"x"}, device_types="cpu")
+        >>> def numpy_sin_inplace(x: Tensor) -> None:
+        >>>     x_np = x.numpy()
+        >>>     np.sin(x_np, out=x_np)
+        >>>
+        >>> x = torch.randn(3)
+        >>> expected = x.sin()
+        >>> numpy_sin_inplace(x)
+        >>> assert torch.allclose(x, expected)
+        >>>
+        >>> # Example of a factory function
+        >>> @torch.library.custom_op("mylib::bar", mutates_args={}, device_types="cpu")
+        >>> def bar(device: torch.device) -> Tensor:
+        >>>     return torch.ones(3)
+        >>>
+        >>> bar("cpu")
+        >>>
+        >>> # Example of a custom op with list inputs
+        >>> @custom_op("mylib::weighted_sum", mutates_args=())
+        >>> def weighted_sum(
+        >>>     tensors: list[Tensor],
+        >>>     weights: list[float],
+        >>> ) -> Tensor:
+        >>>     return sum(t * w for t, w in zip(tensors, weights))
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = torch.randn(3)
+        >>> out = weighted_sum([x, y], [0.3, 0.7])
+
+    """
+
+    def inner(fn: Callable[..., object]) -> CustomOpDef:
+        import torch
+
+        if schema is None:
+            schema_str = torch.library.infer_schema(fn, mutates_args=mutates_args)
+        else:
+            schema_str = schema
+
+        namespace, opname = name.split("::")
+        result = CustomOpDef(namespace, opname, schema_str, fn, tags)
+        if schema is not None:
+            # Check that schema's alias annotations match those of `mutates_args`.
+            expected = set()
+            for arg in result._opoverload._schema.arguments:
+                if arg.alias_info is not None and arg.alias_info.is_write:
+                    expected.add(arg.name)
+            if expected != set(mutates_args):
+                raise ValueError(
+                    f"Attempted to create a custom op with `mutates_args={mutates_args}` "
+                    f"and `schema={schema}. The schema suggests that the op mutates {expected}"
+                    f"which is different from what was provided to us in `mutates_args`. "
+                    f"Please make these consistent."
+                )
+        result.register_kernel(device_types)(fn)
+        return result
+
+    if fn is None:
+        return inner
+    return inner(fn)
+
+
+class CustomOpDef:
+    """CustomOpDef is a wrapper around a function that turns it into a custom op.
+
+    It has various methods for registering additional behavior for this
+    custom op.
+
+    You should not instantiate CustomOpDef directly; instead, use the
+    :func:`torch.library.custom_op` API.
+    """
+
+    def __init__(
+        self,
+        namespace: str,
+        name: str,
+        schema: str,
+        fn: Callable,
+        tags: Sequence[_C.Tag] | None = None,
+    ) -> None:
+        # Fields used to interface with the PyTorch dispatcher
+        self._namespace = namespace
+        self._name = name
+        self._schema = schema
+        self._tags = tags if tags is not None else []
+
+        self._init_fn = fn
+
+        self._backend_fns: dict[str | None, Callable] = {}
+        self._abstract_fn: Callable | None = None
+        self._setup_context_fn: Callable | None = None
+        self._backward_fn: Callable | None = None
+        self._torch_dispatch_fns: dict[type, Callable] = {}
+        self._vmap_fn: Callable | None = None
+        self._autocast_cuda_dtype: _dtype | None = None
+        self._autocast_cpu_dtype: _dtype | None = None
+
+        self._lib = get_library_allowing_overwrite(self._namespace, self._name)
+        self._register_to_dispatcher(self._tags)
+        self._disabled_kernel: set = set()
+        self._used_triton_kernels: list[Any] = list()
+        OPDEFS[self._qualname] = self
+
+    @property
+    def _qualname(self) -> str:
+        return f"{self._namespace}::{self._name}"
+
+    def __repr__(self) -> str:
+        return f"<CustomOpDef({self._qualname})>"
+
+    @contextmanager
+    def set_kernel_enabled(self, device_type: str, enabled: bool = True):
+        """
+        Disable or re-enable an already registered kernel for this custom operator.
+
+        If the kernel is already disabled/enabled, this is a no-op.
+
+        Note:
+            If a kernel is first disabled and then registered, it is disabled until enabled again.
+
+        Args:
+            device_type (str): The device type to disable/enable the kernel for.
+            disable (bool): Whether to disable or enable the kernel.
+
+        Example:
+            >>> inp = torch.randn(1)
+            >>>
+            >>> # define custom op `f`.
+            >>> @custom_op("mylib::f", mutates_args=())
+            >>> def f(x: Tensor) -> Tensor:
+            >>>     return torch.zeros(1)
+            >>>
+            >>> print(f(inp))  # tensor([0.]), default kernel
+            >>>
+            >>> @f.register_kernel("cpu")
+            >>> def _(x):
+            >>>     return torch.ones(1)
+            >>>
+            >>> print(f(inp))  # tensor([1.]), CPU kernel
+            >>>
+            >>> # temporarily disable the CPU kernel
+            >>> with f.set_kernel_enabled("cpu", enabled = False):
+            >>>     print(f(inp))  # tensor([0.]) with CPU kernel disabled
+
+        """
+        action = "enable" if enabled else "disable"
+        originally_disabled = device_type in self._disabled_kernel
+        if device_type not in self._backend_fns:
+            log.warning(
+                "Attempted to %s kernel for %s but no kernel was registered for this device type.",
+                action,
+                device_type,
+            )
+
+        if not enabled:
+            if originally_disabled:
+                log.warning(
+                    "Attempted to disable kernel for %s but it was already disabled.",
+                    device_type,
+                )
+            else:
+                self._disabled_kernel.add(device_type)
+        else:  # enable the kernel
+            if not originally_disabled:
+                log.warning(
+                    "Attempted to enable kernel for  %s but it was already enabled.",
+                    device_type,
+                )
+            else:
+                self._disabled_kernel.remove(device_type)
+
+        try:
+            yield
+        finally:
+            # restore original state
+            if originally_disabled:
+                self._disabled_kernel.add(device_type)
+            else:
+                self._disabled_kernel.discard(device_type)
+
+    def register_kernel(
+        self, device_types: device_types_t, fn: Callable | None = None, /
+    ) -> Callable:
+        """Register an implementation for a device type for this operator.
+
+        Some valid device_types are: "cpu", "cuda", "xla", "mps", "ipu", "xpu".
+        This API may be used as a decorator.
+
+        Args:
+            fn (Callable): The function to register as the implementation for
+                the given device types.
+            device_types (str | Sequence[str]): The device device_types to register an impl to.
+
+        Examples::
+            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+            >>> import torch
+            >>> from torch import Tensor
+            >>> from torch.library import custom_op
+            >>> import numpy as np
+            >>>
+            >>> # Create a custom op that works on cpu
+            >>> @custom_op("mylib::numpy_sin", mutates_args=(), device_types="cpu")
+            >>> def numpy_sin(x: Tensor) -> Tensor:
+            >>>     x_np = x.numpy()
+            >>>     y_np = np.sin(x_np)
+            >>>     return torch.from_numpy(y_np)
+            >>>
+            >>> # Add implementations for the cuda device
+            >>> @numpy_sin.register_kernel("cuda")
+            >>> def _(x):
+            >>>     x_np = x.cpu().numpy()
+            >>>     y_np = np.sin(x_np)
+            >>>     return torch.from_numpy(y_np).to(device=x.device)
+            >>>
+            >>> x_cpu = torch.randn(3)
+            >>> x_cuda = x_cpu.cuda()
+            >>> assert torch.allclose(numpy_sin(x_cpu), x_cpu.sin())
+            >>> assert torch.allclose(numpy_sin(x_cuda), x_cuda.sin())
+
+        """
+
+        def inner(fn):
+            if device_types is None or isinstance(device_types, str):
+                dtypes: list[str | None] = [device_types]
+            else:
+                dtypes = list(device_types)
+            for device_type in dtypes:
+                if device_type not in self._backend_fns:
+
+                    def backend_impl(*args, **kwargs):
+                        result = self._backend_fns[device_type](*args, **kwargs)
+
+                        def get_module():
+                            fn = self._backend_fns[device_type]
+                            return inspect.getmodule(fn)
+
+                        schema = self._opoverload._schema
+                        if not schema._is_view_op():
+                            utils._c_check_aliasing_constraint(
+                                self._name,
+                                args,
+                                kwargs,
+                                result,
+                                get_module,
+                            )
+                        return result
+
+                    if device_type is None:
+                        self._lib.impl(
+                            self._name, backend_impl, "CompositeExplicitAutograd"
+                        )
+                    else:
+                        self._lib.impl(
+                            self._name,
+                            backend_impl,
+                            _C._dispatch_key_for_device(device_type),
+                        )
+
+                # Wrap function to choose between the default implementation or the device-specific
+                # implementation depending on if the kernel is disabled.
+                @torch._disable_dynamo
+                def wrapped_fn(*args, **kwargs):
+                    if device_type in self._disabled_kernel:
+                        return self._init_fn(*args, **kwargs)
+                    else:
+                        return fn(*args, **kwargs)
+
+                self._backend_fns[device_type] = wrapped_fn
+            return fn
+
+        if device_types is not None and not utils.has_tensor_arg(
+            self._opoverload._schema
+        ):
+            device_arg_index = utils.get_device_arg_index(self._opoverload._schema)
+            if device_arg_index is None:
+                raise ValueError(
+                    "Functions without tensor inputs are required to have a `device: torch.device` argument"
+                )
+            self._register_backend_select_dispatcher(device_arg_index)
+
+        # See NOTE: [Supporting decorator and non-decorator usage]
+        if fn is None:
+            return inner
+        return inner(fn)
+
+    def register_fake(self, fn: Callable, /) -> Callable:
+        r"""Register a FakeTensor implementation for this custom op.
+
+        This is necessary to get the operator to work efficiently with torch.compile.
+
+        The Fake impl (sometimes also known as a meta kernel or abstract impl)
+        specifies the behavior of this operator on Tensors that carry no data.
+        Given some input Tensors with certain properties
+        (sizes/strides/storage_offset/device), it specifies what the properties of
+        the output Tensors are.
+
+        Please see :func:`torch.library.register_fake` for more details.
+
+        Args:
+            fn (Callable): The function to register as the FakeTensor
+                implementation.
+
+        Examples:
+            >>> import torch
+            >>> import numpy as np
+            >>> from torch import Tensor
+            >>>
+            >>> # Example 1: an operator without data-dependent output shape
+            >>> @torch.library.custom_op("mylib::linear", mutates_args=())
+            >>> def linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
+            >>>     return (x @ weight.t()) + bias
+            >>>
+            >>> @linear.register_fake
+            >>> def _(x, weight, bias):
+            >>>     assert x.dim() == 2
+            >>>     assert weight.dim() == 2
+            >>>     assert bias.dim() == 1
+            >>>     assert x.shape[1] == weight.shape[1]
+            >>>     assert weight.shape[0] == bias.shape[0]
+            >>>     assert x.device == weight.device
+            >>>     return x.new_empty(x.size(0), weight.size(0))
+            >>>
+            >>> x = torch.randn(2, 2)
+            >>> weight = torch.randn(2, 2)
+            >>> bias = torch.randn(2)
+            >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+            >>> out = torch.compile(linear, fullgraph=True)(x, weight, bias)
+            >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+            >>> assert torch.allclose(out, torch.nn.functional.linear(x, weight, bias))
+            >>>
+            >>> # Example 2: an operator with data-dependent output shape
+            >>> @torch.library.custom_op("mylib::nonzero", mutates_args=())
+            >>> def nonzero(x: Tensor) -> Tensor:
+            >>>     x_np = x.cpu().numpy()
+            >>>     res = np.stack(np.nonzero(x_np), axis=1)
+            >>>     return torch.tensor(res, device=x.device)
+            >>>
+            >>> @nonzero.register_fake
+            >>> def _(x):
+            >>> # Number of nonzero-elements is data-dependent.
+            >>> # Since we cannot peek at the data in an abstract impl,
+            >>> # we use the ctx object to construct a new symint that
+            >>> # represents the data-dependent size.
+            >>>     ctx = torch.library.get_ctx()
+            >>>     nnz = ctx.new_dynamic_size()
+            >>>     shape = [nnz, x.dim()]
+            >>>     result = x.new_empty(shape, dtype=torch.int64)
+            >>>     return result
+            >>>
+            >>> x = torch.tensor([0, 1, 2, 0, 0, 1])
+            >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+            >>> out = torch.compile(nonzero, fullgraph=True)(x)
+            >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+            >>> assert torch.allclose(out, x.nonzero())
+
+        """
+        self._abstract_fn = fn
+        return fn
+
+    def register_effect(self, effect: EffectType | None) -> None:
+        self._lib._register_effectful_op(self._qualname, effect)
+
+    def register_torch_dispatch(
+        self, torch_dispatch_class: Any, fn: Callable | None = None, /
+    ) -> Callable:
+        r"""Registers a torch_dispatch rule for the given operator and ``torch_dispatch_class``.
+
+        This allows for open registration to specify the behavior between the operator
+        and the ``torch_dispatch_class`` without needing to modify the ``torch_dispatch_class``
+        or the operator directly.
+
+        Please see :func:`torch.library.register_torch_dispatch` for examples and more details.
+        """
+
+        def register(fn):
+            if torch_dispatch_class not in self._torch_dispatch_fns:
+
+                def inner(*args, **kwargs):
+                    return self._torch_dispatch_fns[torch_dispatch_class](
+                        *args, **kwargs
+                    )
+
+                self._lib._register_torch_dispatch_rule(
+                    self._name, torch_dispatch_class, inner
+                )
+            self._torch_dispatch_fns[torch_dispatch_class] = fn
+            return fn
+
+        if fn is None:
+            return register
+        else:
+            return register(fn)
+
+    def register_autograd(
+        self,
+        backward: Callable,
+        /,
+        *,
+        setup_context: Callable | None = None,
+    ) -> None:
+        r"""Register a backward formula for this custom op.
+
+        In order for an operator to work with autograd, you need to register
+        a backward formula:
+        1. You must tell us how to compute gradients during the backward pass
+        by providing us a "backward" function.
+        2. If you need any values from the forward to compute gradients, you can
+        use `setup_context` to save values for backward.
+
+        ``backward_fn`` runs during the backward pass. It accepts ``(ctx, *grads)``:
+        - ``grads`` is one or more gradients. The number of gradients matches
+        the number of outputs of the operator.
+        The ``ctx`` object is `the same ctx object <context_method_mixins>`_ used by
+        :class:`torch.autograd.Function`. The semantics of ``backward_fn`` are the
+        same as :meth:`torch.autograd.Function.backward`.
+
+        ``setup_context(ctx, inputs, output)`` runs during the forward pass.
+        Please save quantities needed for backward onto the ``ctx`` object via
+        either :meth:`torch.autograd.function.FunctionCtx.save_for_backward`
+        or assigning them as attributes of ``ctx``. If your custom op has
+        kwarg-only arguments, we expect the signature of ``setup_context``
+        to be ``setup_context(ctx, inputs, keyword_only_inputs, output)``.
+
+        Both ``setup_context_fn`` and ``backward_fn`` must be traceable. That is,
+        they may not directly access :meth:`torch.Tensor.data_ptr` and they must
+        not depend on or mutate global state. If you need a non-traceable backward,
+        you can make it a separate custom_op that you call inside ``backward_fn``.
+
+        If you need different autograd behavior on different devices, then we
+        recommend creating two different custom operators, one for each device
+        that needs different behavior, and switching between them at runtime.
+
+        Examples:
+            >>> import torch
+            >>> import numpy as np
+            >>> from torch import Tensor
+            >>>
+            >>> @torch.library.custom_op("mylib::numpy_sin", mutates_args=())
+            >>> def numpy_sin(x: Tensor) -> Tensor:
+            >>>     x_np = x.cpu().numpy()
+            >>>     y_np = np.sin(x_np)
+            >>>     return torch.from_numpy(y_np).to(device=x.device)
+            >>>
+            >>> def setup_context(ctx, inputs, output) -> Tensor:
+            >>>     x, = inputs
+            >>>     ctx.save_for_backward(x)
+            >>>
+            >>> def backward(ctx, grad):
+            >>>     x, = ctx.saved_tensors
+            >>>     return grad * x.cos()
+            >>>
+            >>> numpy_sin.register_autograd(backward, setup_context=setup_context)
+            >>>
+            >>> x = torch.randn(3, requires_grad=True)
+            >>> y = numpy_sin(x)
+            >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
+            >>> assert torch.allclose(grad_x, x.cos())
+            >>>
+            >>> # Example with a keyword-only arg
+            >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
+            >>> def numpy_mul(x: Tensor, *, val: float) -> Tensor:
+            >>>     x_np = x.cpu().numpy()
+            >>>     y_np = x_np * val
+            >>>     return torch.from_numpy(y_np).to(device=x.device)
+            >>>
+            >>> def setup_context(ctx, inputs, keyword_only_inputs, output) -> Tensor:
+            >>>     ctx.val = keyword_only_inputs["val"]
+            >>>
+            >>> def backward(ctx, grad):
+            >>>     return grad * ctx.val
+            >>>
+            >>> numpy_mul.register_autograd(backward, setup_context=setup_context)
+            >>>
+            >>> x = torch.randn(3, requires_grad=True)
+            >>> y = numpy_mul(x, val=3.14)
+            >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
+            >>> assert torch.allclose(grad_x, torch.full_like(x, 3.14))
+
+        """
+        schema = self._opoverload._schema
+        if not utils.is_functional_schema(schema, allow_valid_view=True):
+            raise RuntimeError(
+                f"Cannot register autograd formula for non-functional operator "
+                f"{self} with schema {schema}. Please create "
+                f"a functional operator and register an autograd formula for that."
+            )
+
+        self._backward_fn = backward
+        self._setup_context_fn = setup_context
+
+    def _register_to_dispatcher(self, tags: Sequence[_C.Tag]) -> None:
+        lib = self._lib
+        schema_str = self._name + self._schema
+        cpp_schema = _C.parse_schema(schema_str)
+        if utils.has_kwarg_only_tensors(cpp_schema):
+            # If you want to support this, the progression is:
+            # - supporting kwarg-only Tensors that are non-differentiable
+            # - supporting kwarg-only Tensors (regardless of differentiability)
+            raise NotImplementedError(
+                f"custom_op with kwarg-only Tensor args. Please make your "
+                f"tensors not kwarg-only. Got: {schema_str}"
+            )
+
+        lib.define(
+            schema_str,
+            tags=[_C.Tag.pt2_compliant_tag, *tags],
+        )
+        self._opoverload = utils.lookup_op(self._qualname)
+
+        def fake_impl(*args, **kwargs):
+            if self._abstract_fn is None:
+                if utils.can_generate_trivial_fake_impl(self._opoverload):
+                    return None
+                raise RuntimeError(
+                    f"There was no fake impl registered for {self}. "
+                    f"This is necessary for torch.compile/export/fx tracing to work. "
+                    f"Please use `{self._init_fn.__name__}.register_fake` to add an "
+                    f"fake impl."
+                )
+            return self._abstract_fn(*args, **kwargs)
+
+        lib._register_fake(self._name, fake_impl, _stacklevel=4)
+
+        autograd_impl = autograd.make_autograd_impl(self._opoverload, self)
+        lib.impl(self._name, autograd_impl, "Autograd", with_keyset=True)
+        schema = self._opoverload._schema
+
+        if schema._is_view_op() or schema.is_mutable:
+            lib.m.register_ad_inplace_or_view_fallback(self._name)  # type: ignore[union-attr]
+
+        if schema.is_mutable:
+            mutated_idxs, mutated_keys = utils.mutated_args_kwargs(schema)
+
+            original_kernel = torch._C._dispatch_get_computed_kernel_for_dispatch_key(
+                f"{lib.ns}::{self._name}", "ADInplaceOrView"
+            )
+
+            def adinplaceorview_impl(keyset, *args, **kwargs):
+                # Handle the mutated idx the user gave us explicitly
+
+                for idx in mutated_idxs:
+                    increment_version(args[idx])
+                for key in mutated_keys:
+                    increment_version(kwargs[key])
+                # Handle view + mutation that are in the schema
+                return original_kernel.call_boxed(keyset, *args, **kwargs)
+
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    message="Warning only once for all operators",
+                    category=UserWarning,
+                )
+                lib.impl(
+                    self._name,
+                    adinplaceorview_impl,
+                    "ADInplaceOrView",
+                    with_keyset=True,
+                )
+
+    def _register_backend_select_dispatcher(self, device_arg_index: int):
+        """
+        Switch on the device argument to select the correct backend to dispatch to.
+        """
+
+        def backend_select(keyset, *args, **kwargs):
+            device = args[device_arg_index].type
+            if device not in self._backend_fns:
+                raise RuntimeError(
+                    f"{self._name} does not have a kernel registered for {device}. "
+                    "Please use register_kernel to do so."
+                )
+            # pyrefly: ignore [bad-argument-type]
+            dispatch_key = _C._dispatch_key_for_device(device)
+            dispatch_key = getattr(_C.DispatchKey, dispatch_key)
+            return self._opoverload.redispatch(
+                _C.DispatchKeySet(dispatch_key), *args, **kwargs
+            )
+
+        self._lib.impl(self._name, backend_select, "BackendSelect", with_keyset=True)
+
+    def __call__(self, *args, **kwargs):
+        return self._opoverload(*args, **kwargs)
+
+    def register_vmap(
+        self,
+        func: Callable | None = None,
+    ):
+        r"""Register a vmap implementation to support :func:`torch.vmap` for this custom op.
+
+        This API may be used as a decorator.
+
+        In order for an operator to work with :func:`torch.vmap`, you may need to register a
+        vmap implementation in the following signature:
+
+            ``vmap_func(info, in_dims: Tuple[Optional[int]], *args, **kwargs)``,
+
+        where ``*args`` and ``**kwargs`` are the arguments and kwargs for ``op``.
+
+        It specifies how do we compute the batched version of ``op`` given inputs with an additional
+        dimension (specified by ``in_dims``).
+
+        For each arg in ``args``, ``in_dims`` has a corresponding ``Optional[int]``. It is ``None``
+        if the arg is not a Tensor or if the arg is not being vmapped over, otherwise, it is an integer
+        specifying what dimension of the Tensor is being vmapped over.
+
+        ``info`` is a collection of additional metadata that may be helpful:
+        ``info.batch_size`` specifies the size of the dimension being vmapped over, while
+        ``info.randomness`` is the ``randomness`` option that was passed to :func:`torch.vmap`.
+
+        The return of the function ``func`` is a tuple of ``(output, out_dims)``. Similar to ``in_dims``,
+        ``out_dims`` should be of the same structure as ``output`` and contain one ``out_dim``
+        per output that specifies if the output has the vmapped dimension and what index it is in.
+
+        Examples:
+            >>> import torch
+            >>> import numpy as np
+            >>> from torch import Tensor
+            >>> from typing import Tuple
+            >>>
+            >>> def to_numpy(tensor):
+            >>>     return tensor.cpu().numpy()
+            >>>
+            >>> lib = torch.library.Library("mylib", "FRAGMENT")
+            >>> @torch.library.custom_op("mylib::numpy_cube", mutates_args=())
+            >>> def numpy_cube(x: Tensor) -> Tuple[Tensor, Tensor]:
+            >>>     x_np = to_numpy(x)
+            >>>     dx = torch.tensor(3 * x_np ** 2, device=x.device)
+            >>>     return torch.tensor(x_np ** 3, device=x.device), dx
+            >>>
+            >>> def numpy_cube_vmap(info, in_dims, x):
+            >>>     result = numpy_cube(x)
+            >>>     return result, (in_dims[0], in_dims[0])
+            >>>
+            >>> numpy_cube.register_vmap(numpy_cube_vmap)
+            >>>
+            >>> x = torch.randn(3)
+            >>> torch.vmap(numpy_cube)(x)
+            >>>
+            >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
+            >>> def numpy_mul(x: Tensor, y: Tensor) -> Tensor:
+            >>>     return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
+            >>>
+            >>> @numpy_mul.register_vmap
+            >>> def numpy_mul_vmap(info, in_dims, x, y):
+            >>>     x_bdim, y_bdim = in_dims
+            >>>     x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
+            >>>     y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
+            >>>     result = x * y
+            >>>     result = result.movedim(-1, 0)
+            >>>     return result, 0
+            >>>
+            >>>
+            >>> x = torch.randn(3)
+            >>> y = torch.randn(3)
+            >>> torch.vmap(numpy_mul)(x, y)
+        """
+        from torch._functorch.autograd_function import custom_function_call_vmap_helper
+        from torch._functorch.pyfunctorch import retrieve_current_functorch_interpreter
+
+        def register(func):
+            need_register = self._vmap_fn is None
+            self._vmap_fn = func
+
+            if need_register:
+
+                def wrapped_func(keyset, *args, **kwargs):
+                    interpreter = retrieve_current_functorch_interpreter()
+                    return custom_function_call_vmap_helper(
+                        # pyrefly: ignore[bad-argument-type]
+                        interpreter,
+                        # pyrefly: ignore[bad-argument-type]
+                        self._vmap_fn,
+                        self._opoverload,
+                        *args,
+                        **kwargs,
+                    )
+
+                self._lib.impl(
+                    self._name, wrapped_func, "FuncTorchBatched", with_keyset=True
+                )
+
+        if func is None:
+            return register
+        else:
+            return register(func)
+
+    def register_autocast(
+        self,
+        device_type: str,
+        cast_inputs: _dtype,
+    ):
+        r"""Register an autocast dispatch rule for this custom op.
+
+        Valid `device_type` include: "cpu" and "cuda".
+
+        Args:
+            op (str | OpOverload): The operator to register an autocast dispatch rule to.
+            device_type(str):  Device type to use. 'cuda' or 'cpu'.
+                The type is the same as the `type` attribute of a :class:`torch.device`.
+                Thus, you may obtain the device type of a tensor using `Tensor.device.type`.
+            cast_inputs (:class:`torch.dtype`): When custom op runs in an autocast-enabled region,
+                casts incoming floating-point Tensors to the target dtype (non-floating-point Tensors
+                are not affected), then executes custom op with autocast disabled.
+            lib (Optional[Library]): If provided, the lifetime of this registration
+
+        Examples::
+            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+            >>> import torch
+            >>> from torch import Tensor
+            >>> from torch.library import custom_op
+            >>>
+            >>> # Create a custom op that works on cuda
+            >>> @torch.library.custom_op("mylib::my_sin", mutates_args=())
+            >>> def my_sin(x: Tensor) -> Tensor:
+            >>>     return torch.sin(x)
+            >>>
+            >>> # Register autocast dispatch rule for the cuda device
+            >>> torch.library.register_autocast("mylib::my_sin", "cuda", torch.float16)
+            >>>
+            >>> x = torch.randn(3, dtype=torch.float32, device="cuda")
+            >>> with torch.autocast("cuda", dtype=torch.float16):
+            >>>     y = torch.ops.mylib.my_sin(x)
+            >>> assert y.dtype == torch.float16
+
+        """
+        if not isinstance(device_type, str):
+            raise ValueError(
+                f"Expected `device_type` of type `str`, got: `{type(device_type)}`"
+            )
+        if device_type not in ["cpu", "cuda"]:
+            raise ValueError(f"Unknown device type: {device_type}")
+
+        need_register_cuda = self._autocast_cuda_dtype is None
+        need_register_cpu = self._autocast_cpu_dtype is None
+        if device_type == "cuda":
+            self._autocast_cuda_dtype = cast_inputs
+        else:
+            self._autocast_cpu_dtype = cast_inputs
+
+        def kernel(_, *args, **kwargs):
+            if len(kwargs) != 0:
+                raise AssertionError(
+                    f"Custom ops do not support kwargs yet, got {list(kwargs.keys())}"
+                )
+            autocast_keyset = torch._C.DispatchKeySet(
+                torch._C.DispatchKey.AutocastCPU
+            ) | torch._C.DispatchKeySet(torch._C.DispatchKey.AutocastCUDA)
+            with torch._C._ExcludeDispatchKeyGuard(autocast_keyset):
+                return self._opoverload(*_cast(args, device_type, cast_inputs))
+
+        if need_register_cuda and self._autocast_cuda_dtype:
+            self._lib.impl(self._name, kernel, "AutocastCUDA", with_keyset=True)
+        elif need_register_cpu and self._autocast_cpu_dtype:
+            self._lib.impl(self._name, kernel, "AutocastCPU", with_keyset=True)
+
+        return kernel
+
+
+# TODO: Merge this function with torch.amp.autocast_mode._cast, and refactor it
+# into a utility function once custom ops support arbitrary input types.
+def _cast(value, device_type: str, dtype: _dtype):
+    if isinstance(value, torch.Tensor):
+        is_eligible = (
+            value.is_floating_point()
+            and value.device.type == device_type
+            and (value.dtype is not torch.float64)
+        )
+        return value.to(dtype) if is_eligible else value
+    elif isinstance(value, (str, bytes)):
+        return value
+    elif isinstance(value, collections.abc.Iterable):
+        iterable = (_cast(v, device_type, dtype) for v in value)
+        if isinstance(value, (list, tuple)):
+            return type(value)(iterable)
+        else:
+            return iterable
+    else:
+        return value
+
+
+def increment_version(val: Any) -> None:
+    if isinstance(val, Tensor):
+        torch.autograd.graph.increment_version(val)
+    elif isinstance(val, (tuple, list)):
+        for v in val:
+            if isinstance(v, Tensor):
+                torch.autograd.graph.increment_version(v)
+
+
+# NOTE: [Supporting decorator and non-decorator usage]
+#
+# Some APIs may be both used as a decorator and not as a decorator.
+# For example:
+#
+# >>> def fn(x):
+# >>>     return x.sin()
+# >>>
+# >>> # Usage 1: not as a decorator
+# >>> numpy_sin.register_kernel("cuda", fn)
+# >>>
+# >>> # Usage 2: as a decorator
+# >>> @numpy_sin.register_kernel("cuda")
+# >>> def fn2(x):
+# >>>     return x.sin
+#
+# The way we support this is that `register_kernel` accepts an optional `fn`.
+# If `fn` is provided (Usage 1), then we know that the user is using it not
+# as a decorator.
+# If `fn` is not provided (Usage 2), then `register_kernel` needs to return a
+# decorator.
+
+
+OPDEF_TO_LIB: dict[str, "torch.library.Library"] = {}
+OPDEFS: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
+
+
+def get_library_allowing_overwrite(
+    namespace: str, name: str
+) -> "torch.library.Library":
+    qualname = f"{namespace}::{name}"
+
+    if qualname in OPDEF_TO_LIB:
+        OPDEF_TO_LIB[qualname]._destroy()
+        del OPDEF_TO_LIB[qualname]
+
+    lib = torch.library.Library(namespace, "FRAGMENT")  # noqa: TOR901
+    OPDEF_TO_LIB[qualname] = lib
+    return lib
+
+
+def _maybe_get_opdef(
+    op: CustomOpDef | _ops.OpOverload | str,
+) -> CustomOpDef | None:
+    if isinstance(op, CustomOpDef):
+        return op
+    if isinstance(op, _ops.OpOverload):
+        op = op._name
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str, got {type(op)}")
+    if op in OPDEFS:
+        return OPDEFS[op]
+    return None

kronos-venv/lib/python3.12/site-packages/torch/_library/effects.py

@@ -0,0 +1,84 @@
+from enum import Enum
+
+import torch
+
+
+class EffectType(Enum):
+    ORDERED = "Ordered"
+
+
+from torch._library.utils import RegistrationHandle
+
+
+# These classes do not have side effects as they just store quantization
+# params, so we dont need to mark them as ordered
+skip_classes = (
+    "__torch__.torch.classes.quantized.Conv2dPackedParamsBase",
+    "__torch__.torch.classes.quantized.Conv3dPackedParamsBase",
+    "__torch__.torch.classes.quantized.EmbeddingPackedParamsBase",
+    "__torch__.torch.classes.quantized.LinearPackedParamsBase",
+    "__torch__.torch.classes.xnnpack.Conv2dOpContext",
+    "__torch__.torch.classes.xnnpack.LinearOpContext",
+    "__torch__.torch.classes.xnnpack.TransposeConv2dOpContext",
+)
+
+
+class EffectHolder:
+    """A holder where one can register an effect impl to."""
+
+    def __init__(self, qualname: str):
+        self.qualname: str = qualname
+        self._set_default_effect()
+
+    def _set_default_effect(self) -> None:
+        self._effect: EffectType | None = None
+
+        # If the op contains a ScriptObject input, we want to mark it as having effects
+        namespace, opname = torch._library.utils.parse_namespace(self.qualname)
+        split = opname.split(".")
+        if len(split) > 1:
+            if len(split) != 2:
+                raise AssertionError(
+                    f"Tried to split {opname} based on '.' but found more than 1 '.'"
+                )
+            opname, overload = split
+        else:
+            overload = ""
+
+        if namespace == "higher_order":
+            return
+
+        opname = f"{namespace}::{opname}"
+        if torch._C._get_operation_overload(opname, overload) is not None:
+            # Since we call this when destroying the library, sometimes the
+            # schema will be gone already at that time.
+            schema = torch._C._get_schema(opname, overload)
+            for arg in schema.arguments:
+                if isinstance(arg.type, torch.ClassType):
+                    type_str = arg.type.str()  # pyrefly: ignore[missing-attribute]
+                    if type_str in skip_classes:
+                        continue
+                    self._effect = EffectType.ORDERED
+                    return
+
+    @property
+    def effect(self) -> EffectType | None:
+        return self._effect
+
+    @effect.setter
+    def effect(self, _):
+        raise RuntimeError("Unable to directly set kernel.")
+
+    def register(self, effect: EffectType | None) -> RegistrationHandle:
+        """Register an effect
+
+        Returns a RegistrationHandle that one can use to de-register this
+        effect.
+        """
+        self._effect = effect
+
+        def deregister_effect():
+            self._set_default_effect()
+
+        handle = RegistrationHandle(deregister_effect)
+        return handle

kronos-venv/lib/python3.12/site-packages/torch/_library/fake_class_registry.py

@@ -0,0 +1,510 @@
+# mypy: allow-untyped-defs
+import copy
+import logging
+from typing import Any, Protocol
+
+import torch
+from torch._library.utils import parse_namespace
+from torch.utils._python_dispatch import _disable_current_modes
+
+
+log = logging.getLogger(__name__)
+
+
+class FakeScriptObject:
+    def __init__(
+        self, wrapped_obj: Any, script_class_name: str, x: torch.ScriptObject | None
+    ):
+        # Use object.__setattr__ to bypass our custom __setattr__ during initialization
+        object.__setattr__(self, "wrapped_obj", wrapped_obj)
+        object.__setattr__(self, "script_class_name", script_class_name)
+
+        from torch._library.opaque_object import is_opaque_type
+
+        # We dont want to deepcopy when tracing with opaque objects because
+        # if a mutation happens intentionally (Ex. caching in device mesh)
+        # then we want it to be recorded on the real object
+        real_obj = x
+        if not is_opaque_type(type(x)):
+            try:
+                with _disable_current_modes():
+                    real_obj = copy.deepcopy(x)
+            except (RuntimeError, TypeError) as e:
+                log.warning(  # noqa: G200
+                    "Unable to deepcopy the custom object %s due to %s. "
+                    "Defaulting to the user given object. This might be "
+                    "dangerous as side effects may be directly applied "
+                    "to the object.",
+                    script_class_name,
+                    e,
+                )
+
+        object.__setattr__(self, "real_obj", real_obj)
+
+    def __getattribute__(self, name):
+        try:
+            return super().__getattribute__(name)
+        except AttributeError as e:
+            raise AttributeError(
+                f"Tried to call __getattr__ with attr '{name}' on a FakeScriptObject, "
+                "implying that you are calling this inside of a fake kernel. "
+                "The fake kernel should not depend on the contents of the "
+                "OpaqueObject at all, so we're erroring out. If this attr is "
+                "a method or constant attribute, you can allow this member access by "
+                "registering it via `register_opaque_type(members=...)`."
+            ) from e
+
+    def __setattr__(self, name, value):
+        raise AttributeError(
+            f"Tried to call __setattr__ with attr '{name}' on a FakeScriptObject, "
+            "implying that you are calling this inside of a fake kernel. "
+            "The fake kernel should not depend on the contents of the "
+            "OpaqueObject at all, so we're erroring out. If you need this"
+            "functionality, consider creating a custom TorchBind Object instead"
+            "(but note that this is more difficult)."
+        )
+
+    def __getitem__(self, key):
+        # This is needed for DeviceMesh support
+        return self.real_obj[key]
+
+    def __eq__(self, other):
+        if self is other:
+            return True
+        # Get real_obj without triggering custom __getattribute__
+        self_real = object.__getattribute__(self, "real_obj")
+        if isinstance(other, FakeScriptObject):
+            other_real = object.__getattribute__(other, "real_obj")
+            # For reference types, identity check first
+            if self_real is other_real:
+                return True
+            # Fall back to equality check
+            return self_real == other_real
+        # Compare with the real object directly
+        return self_real == other
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __hash__(self):
+        # Use real_obj's hash if available, otherwise use object id
+        real_obj = object.__getattribute__(self, "real_obj")
+        try:
+            return hash(real_obj)
+        except TypeError:
+            # Object is not hashable, use identity-based hash
+            return id(real_obj)
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> "FakeScriptObject":
+        if id(self) in memo:
+            return memo[id(self)]
+        new_obj = FakeScriptObject.__new__(FakeScriptObject)
+        memo[id(self)] = new_obj
+        object.__setattr__(
+            new_obj, "wrapped_obj", copy.deepcopy(self.wrapped_obj, memo)
+        )
+        object.__setattr__(new_obj, "script_class_name", self.script_class_name)
+        # Disable dispatch modes during deepcopy of real_obj and attribute
+        # access to prevent tensor operations (e.g. storage cloning, property
+        # access on DeviceMesh) from going through proxy tracing or
+        # functionalization.
+        with _disable_current_modes():
+            new_real_obj = copy.deepcopy(self.real_obj, memo)
+            object.__setattr__(new_obj, "real_obj", new_real_obj)
+            for name, value in self.__dict__.items():
+                if name not in ("wrapped_obj", "script_class_name", "real_obj"):
+                    if isinstance(value, FakeScriptMethod):
+                        object.__setattr__(
+                            new_obj,
+                            name,
+                            FakeScriptMethod(new_obj, value.method_name, value.schema),
+                        )
+                    else:
+                        if hasattr(new_real_obj, name):
+                            object.__setattr__(
+                                new_obj, name, getattr(new_real_obj, name)
+                            )
+                        else:
+                            object.__setattr__(new_obj, name, value)
+        return new_obj
+
+
+def maybe_unwrap_fake_script_object(obj: Any) -> Any:
+    """If obj is a FakeScriptObject, return the underlying real object."""
+    if isinstance(obj, FakeScriptObject):
+        return obj.real_obj
+    return obj
+
+
+class FakeScriptMethod:
+    def __init__(
+        self,
+        self_fake_obj: FakeScriptObject,
+        method_name: str,
+        schema: torch.FunctionSchema | None,
+    ):
+        self.self_fake_obj = self_fake_obj
+        self.method_name = method_name
+        self.schema = schema
+
+    def __call__(self, *args, **kwargs):
+        from torch._higher_order_ops.torchbind import call_torchbind
+
+        return call_torchbind(self.self_fake_obj, self.method_name, *args, **kwargs)
+
+
+class HasStaticMethodFromReal(Protocol):
+    @classmethod
+    def from_real(cls, real_obj: torch.ScriptObject):
+        pass
+
+
+class FakeClassRegistry:
+    def __init__(self) -> None:
+        self._registered_class: dict[str, Any] = {}
+
+    def has_impl(self, full_qualname: str) -> bool:
+        return full_qualname in self._registered_class
+
+    def get_impl(self, full_qualname: str) -> Any:
+        self._check_registered(full_qualname)
+        return self._registered_class[full_qualname]
+
+    def register(self, full_qualname: str, fake_class=None) -> None:
+        if self.has_impl(full_qualname):
+            log.warning(
+                "%s is already registered. Previous fake class is overridden with  %s.",
+                full_qualname,
+                fake_class,
+            )
+        self._registered_class[full_qualname] = fake_class
+
+    def deregister(self, full_qualname: str) -> Any:
+        if not self.has_impl(full_qualname):
+            log.warning(
+                "Cannot deregister %s. Please use register_fake_class to register it first."
+                " Or do you dereigster it twice?",
+                full_qualname,
+            )
+        else:
+            return self._registered_class.pop(full_qualname)
+
+    def clear(self) -> None:
+        self._registered_class.clear()
+
+    def _check_registered(self, full_qualname: str) -> None:
+        if full_qualname not in self._registered_class:
+            raise RuntimeError(
+                f"{full_qualname} is not registered. Please use register_fake_class to register it first."
+            )
+
+
+global_fake_class_registry = FakeClassRegistry()
+
+
+# TODO: add this check at compile time for __obj_flatten__.
+def _check_valid_flat_script_obj(flat_x):
+    if not isinstance(flat_x, tuple):
+        raise RuntimeError("Expect flat x to be a tuple.")
+
+    for tp in flat_x:
+        if not isinstance(tp, tuple):
+            raise RuntimeError("Expect flat x to be a tuple of tuples.")
+
+        if not len(tp) == 2 or not isinstance(tp[0], str):
+            raise RuntimeError(
+                "Expect element of flat x to be a tuple of two elements with first element being a string"
+            )
+
+
+def tracing_with_real(x: torch.ScriptObject) -> bool:
+    if not hasattr(x, "tracing_mode"):
+        return False
+
+    if x.tracing_mode() not in ["real", "fake"]:
+        raise AssertionError(
+            f"tracing_mode can be either real or fake but got {x.tracing_mode()}"
+        )
+    return x.tracing_mode() == "real"
+
+
+def maybe_to_fake_obj(
+    fake_mode,
+    x: Any,
+) -> FakeScriptObject | torch.ScriptObject:
+    import torch.utils._pytree as pytree
+
+    # When tracing with real mode, people should implement meta kernels that can
+    # handle the case of real script object + fake tensor inputs.
+    if tracing_with_real(x):
+        return x
+
+    from torch._library.opaque_object import (
+        FakeOpaqueObject,
+        get_opaque_obj_info,
+        get_opaque_type_name,
+        is_opaque_type,
+        OpaqueTypeStr,
+    )
+
+    x_type = type(x)
+    if is_opaque_type(x_type):
+        type_name = OpaqueTypeStr if x is None else get_opaque_type_name(x_type)
+        fake_x_wrapped = FakeScriptObject(FakeOpaqueObject(), type_name, x)
+
+        # Set specified members onto the fake object
+        opaque_info = get_opaque_obj_info(x_type)
+        if opaque_info is None:
+            raise AssertionError(f"opaque_info for type {x_type} must not be None")
+        for attr_name in opaque_info.members:
+            with _disable_current_modes():
+                if not hasattr(x, attr_name):
+                    raise TypeError(
+                        f"Opaque object of type '{type_name}' was specified to have member "
+                        f"'{attr_name}', but this doesn't actually exist in the object."
+                    )
+                object.__setattr__(fake_x_wrapped, attr_name, getattr(x, attr_name))
+
+        return fake_x_wrapped
+    else:
+        # x.__obj_flatten__() could be calling some tensor operations inside but we don't
+        # want to call these ops in surrounding dispatch modes when executing it.
+        # Otherwise, for example, the fake tensor modes will error out when the tensors inside
+        # script object execute some operations like clone if allow_non_fake_input flag is set.
+        with _disable_current_modes():
+            flat_x = x.__obj_flatten__()  # type: ignore[attr-defined]
+
+        _check_valid_flat_script_obj(flat_x)
+
+        with fake_mode:
+            from torch._higher_order_ops.utils import _tensor_storage
+
+            storage_map = {
+                _tensor_storage(inp): i
+                for i, inp in enumerate(flat_x)
+                if isinstance(inp, torch.Tensor)
+            }
+            alias_map = {
+                i: storage_map[_tensor_storage(inp)]
+                for i, inp in enumerate(flat_x)
+                if isinstance(inp, torch.Tensor)
+                and storage_map[_tensor_storage(inp)] != i
+            }
+            if len(alias_map) > 0:
+                log.warning(
+                    "Detected script object %s has aliasing relationship among its tensors. "
+                    "Flattened obj: %s. Aliasing tensor indices: %s. "
+                    "This is not supported and may cause unexpected behavior.",
+                    x,
+                    flat_x,
+                    alias_map,
+                )
+
+            # This breaks the aliasing relationship among the tensors inside the torchbind object
+            # This is bad but since we don't need to preserve the aliasing relationship anyway and
+            # we state clearly that aliasing relationship is not preserved in the doc so this might be OK.
+            fake_flattened = pytree.tree_map_only(
+                torch.Tensor,
+                lambda t: torch.empty_strided(
+                    t.size(),
+                    t.stride(),
+                    device=t.device,
+                    dtype=t.dtype,
+                    requires_grad=t.requires_grad,
+                    layout=t.layout,
+                ),
+                flat_x,
+            )
+
+        fake_x = _find_fake_class_for_script_object(x).__obj_unflatten__(fake_flattened)
+
+    fake_x_wrapped = FakeScriptObject(fake_x, x._type().qualified_name(), x)  # type: ignore[attr-defined]
+
+    for name in x._method_names():  # type: ignore[attr-defined]
+        attr = getattr(fake_x, name, None)
+        if attr is not None:
+            if not callable(attr):
+                raise RuntimeError(f"Expect {name} to be a callable but got {attr}.")
+
+            real_attr = getattr(x, name)  # type: ignore[attr-defined]
+
+            # real attr sometimes is not torch.ScriptMethod thus doesn't have schema e.g. __init___ or __eq__
+            method_schema: torch.FunctionSchema | None = None
+            if isinstance(real_attr, torch.ScriptMethod):
+                method_schema = real_attr.schema  # type: ignore[attr-defined]
+
+            # Bypasses our custom setattr function
+            object.__setattr__(
+                fake_x_wrapped,
+                name,
+                FakeScriptMethod(fake_x_wrapped, name, method_schema),
+            )
+        else:
+            override_skip_list = {"__obj_flatten__", "__getstate__", "__setstate__"}
+            if name not in override_skip_list:
+                log.warning("fake object of %s doesn't implement method %s.", x, name)
+    return fake_x_wrapped
+
+
+def register_fake_class(qualname, fake_class: HasStaticMethodFromReal | None = None):
+    r"""Register a fake implementation for this class.
+
+    It's in the same spirit of registering a fake implementation for
+    an operator but with the difference that it
+    associates a fake class with the original torch bind class (registered
+    with torch::class_). In this way, torch.compile can handle them properly
+    in components such as Dynamo and AOTAutograd.
+
+    This API may be used as a decorator (see example). For the fake class, users
+    are required to provide a from_real classmethod that takes a real object and
+    returns an instance of the fake class. All tensors in the fake object should also
+    be properly fakified with to_fake_tensor() in from_real.
+
+
+    Examples:
+        # For a custom class Foo defined in test_custom_class_registration.cpp:
+
+        TORCH_LIBRARY(_TorchScriptTesting, m) {
+          m.class_<TensorQueue>("_TensorQueue")
+            .def(torch::init<at::Tensor>())
+            .def("push", &TensorQueue::push)
+            .def("pop", &TensorQueue::pop)
+            .def("top", &TensorQueue::top)
+            .def("size", &TensorQueue::size)
+            .def("clone_queue", &TensorQueue::clone_queue)
+            .def("__obj_flatten__", &TensorQueue::__obj_flatten__)
+            .def_pickle(
+                // __getstate__
+                [](const c10::intrusive_ptr<TensorQueue>& self)
+                    -> c10::Dict<std::string, at::Tensor> {
+                  return self->serialize();
+                },
+                // __setstate__
+                [](c10::Dict<std::string, at::Tensor> data)
+                    -> c10::intrusive_ptr<TensorQueue> {
+                  return c10::make_intrusive<TensorQueue>(std::move(data));
+                });
+            };
+        # We could register a fake class FakeTensorQueue in Python as follows:
+        import torch
+
+        @torch._library.register_fake_class("_TorchScriptTesting::_TensorQueue")
+        class FakeTensorQueue:
+            def __init__(self, queue):
+                self.queue = queue
+
+            @classmethod
+            def __obj_unflatten__(cls, flattened_ctx):
+                return cls(**dict(ctx))
+
+            def push(self, x):
+                self.queue.append(x)
+
+            def pop(self):
+                return self.queue.pop(0)
+
+            def size(self):
+                return len(self.queue)
+
+    In this example, the original TensorQeue need to add a __obj_flatten__ method
+    to the class TensorQueue and the flattened result is passed into FakeTensorQueue's
+    __obj_unflatten__ as inputs to create a fake class. This protocol allows pytorch to look
+    at the contents of the script object and properly handle them in the subsystems
+    like dynamo, aot_aotugrad or more.
+    """
+
+    def inner(fake_class: HasStaticMethodFromReal):
+        ns, name = parse_namespace(qualname)
+
+        # This also checks whether the referred torch::class_ exists.
+        torch._C._get_custom_class_python_wrapper(ns, name)
+
+        from_method = getattr(fake_class, _CONVERT_FROM_REAL_NAME, None)
+        if not from_method:
+            raise RuntimeError(
+                f"{fake_class} doesn't define a classmethod {_CONVERT_FROM_REAL_NAME}."
+            )
+
+        if not isinstance(fake_class.__dict__[_CONVERT_FROM_REAL_NAME], classmethod):
+            raise RuntimeError(
+                f"{_CONVERT_FROM_REAL_NAME} method is not a classmethod."
+            )
+
+        global_fake_class_registry.register(_full_qual_class_name(qualname), fake_class)
+        return fake_class
+
+    if fake_class is None:
+        return inner
+    return inner(fake_class)
+
+
+def deregister_fake_class(qualname):
+    return global_fake_class_registry.deregister(_full_qual_class_name(qualname))
+
+
+def has_fake_class(full_qualname) -> bool:
+    return global_fake_class_registry.has_impl(full_qualname)
+
+
+def find_fake_class(full_qualname) -> Any | None:
+    if not has_fake_class(full_qualname):
+        return None
+    return global_fake_class_registry.get_impl(full_qualname)
+
+
+def _full_qual_class_name(qualname: str) -> str:
+    ns, name = parse_namespace(qualname)
+    return "__torch__.torch.classes." + ns + "." + name
+
+
+def _is_script_object(obj: Any) -> bool:
+    return isinstance(
+        obj, torch.ScriptObject
+    ) and obj._type().qualified_name().startswith(  # type: ignore[attr-defined]
+        "__torch__.torch.classes"
+    )
+
+
+# Return the namespace and class name from fully qualified name.
+def _ns_and_class_name(full_qualname: str) -> tuple[str, str]:
+    splits = full_qualname.split(".")
+    if len(splits) != 5:
+        raise AssertionError(f"Could not split {full_qualname=}, expected 5 parts")
+    _torch, _torch_ns, _classes, ns, class_name = splits
+    return ns, class_name
+
+
+def _find_fake_class_for_script_object(x: torch.ScriptObject) -> Any:
+    full_qualname = x._type().qualified_name()  # type: ignore[attr-defined]
+    ns, class_name = _ns_and_class_name(full_qualname)
+    fake_class = find_fake_class(full_qualname)
+    if fake_class is None:
+        raise RuntimeError(
+            f" ScriptObject's {full_qualname} haven't registered a fake class."
+            f" Please use register_fake_class({ns}::{class_name}) to annotate a fake class for the script obj."
+            f" Specifically, create a python class that implements a fake version for all the methods"
+            f" that're used in the program and put annotated class in the program e.g. after loading the library."
+            f" The fake methods can be written in the same way as a meta kernel for an operator but need to additionally"
+            f" simulate the object's states. Be sure to add a {_CONVERT_FROM_REAL_NAME} classmethod"
+            f" to enable creating a fake obj from a real one."
+        )
+    return fake_class
+
+
+_CONVERT_FROM_REAL_NAME = "__obj_unflatten__"
+
+
+def _fake_obj_from_real(fake_mode, x) -> Any:
+    fake_class = _find_fake_class_for_script_object(x)
+
+    from_real_method = getattr(fake_class, _CONVERT_FROM_REAL_NAME, None)
+    if not from_real_method:
+        raise RuntimeError(
+            f"{fake_class} must define a classmethod {_CONVERT_FROM_REAL_NAME}"
+            f" that converts the real object to the fake object."
+        )
+
+    # from_real defined by user need the ctx to fakify the tensor states.
+    ctx = torch._library.fake_impl.FakeImplCtx(fake_mode, None)
+    with torch._library.fake_impl.set_ctx_getter(lambda: ctx):
+        return fake_class.from_real(x)

kronos-venv/lib/python3.12/site-packages/torch/_library/fake_impl.py

@@ -0,0 +1,229 @@
+# mypy: allow-untyped-defs
+import contextlib
+import functools
+from collections.abc import Callable
+from typing_extensions import deprecated
+
+import torch
+from torch._library.utils import Kernel, RegistrationHandle
+
+
+class FakeImplHolder:
+    """A holder where one can register an fake impl to."""
+
+    def __init__(self, qualname: str):
+        self.qualname: str = qualname
+        # kernels stores all registered fake kernels, ordered by registration
+        # time ascendingly (newer registration after older registration). If an
+        # operator library gets loaded that overrides an existing fake kernel,
+        # both kernels will be in the list, but the newest one will be the one
+        # that is run. If the library is unloaded, we will remove the kernel
+        # from this list.
+        self.kernels: list[Kernel] = []
+
+    @property
+    def kernel(self):
+        if len(self.kernels) == 0:
+            return None
+        return self.kernels[-1]
+
+    @kernel.setter
+    def kernel(self, value):
+        raise RuntimeError("Unable to directly set kernel.")
+
+    def register(
+        self, func: Callable, source: str, lib, *, allow_override=False
+    ) -> RegistrationHandle:
+        """Register an fake impl.
+
+        Returns a RegistrationHandle that one can use to de-register this
+        fake impl.
+        """
+
+        if not allow_override:
+            if self.kernel is not None:
+                raise RuntimeError(
+                    f"register_fake(...): the operator {self.qualname} "
+                    f"already has an fake impl registered at "
+                    f"{self.kernel.source}."
+                )
+            if torch._C._dispatch_has_kernel_for_dispatch_key(self.qualname, "Meta"):
+                raise RuntimeError(
+                    f"register_fake(...): the operator {self.qualname} "
+                    f"already has an DispatchKey::Meta implementation via a "
+                    f"pre-existing torch.library or TORCH_LIBRARY registration. "
+                    f"Please either remove that registration or don't call "
+                    f"register_fake."
+                )
+
+            if torch._C._dispatch_has_kernel_for_dispatch_key(
+                self.qualname, "CompositeImplicitAutograd"
+            ):
+                raise RuntimeError(
+                    f"register_fake(...): the operator {self.qualname} "
+                    f"already has an implementation for this device type via a "
+                    f"pre-existing registration to "
+                    f"DispatchKey::CompositeImplicitAutograd."
+                    f"CompositeImplicitAutograd operators do not need an fake "
+                    f"impl; "
+                    f"instead, the operator will decompose into its constituents "
+                    f"and those "
+                    f"can have fake impls defined on them."
+                )
+
+        # Store the kernel in this holder
+        kernel = Kernel(func, source)
+        self.kernels.append(kernel)
+
+        def deregister_fake_kernel():
+            self.kernels.remove(kernel)
+
+        meta_kernel = construct_meta_kernel(self.qualname, self)
+        lib.impl(self.qualname, meta_kernel, "Meta", allow_override=allow_override)
+
+        handle = RegistrationHandle(deregister_fake_kernel)
+        return handle
+
+
+def construct_meta_kernel(qualname: str, fake_impl_holder: FakeImplHolder) -> Callable:
+    if fake_impl_holder.kernel is None:
+        raise AssertionError("fake_impl_holder.kernel must not be None")
+
+    @functools.wraps(fake_impl_holder.kernel.func)
+    def meta_kernel(*args, **kwargs):
+        if fake_impl_holder.kernel is None:
+            raise AssertionError("fake_impl_holder.kernel must not be None")
+        source = fake_impl_holder.kernel.source
+
+        def error_on_ctx():
+            raise RuntimeError(
+                f"{qualname} ({source}): You're trying to run this operator "
+                f"with meta Tensors (as opposed to FakeTensors), but this "
+                f"operator may return an output Tensor with data-dependent shape. Meta "
+                f"Tensors don't support operators with outputs that have data-dependent shapes "
+                f"but FakeTensors do. "
+                f"If your operator does not return an output with data-dependent shape, "
+                f"make sure the FakeTensor and/or meta kernel does not call "
+                f"torch.library.get_ctx(). Otherwise, please use FakeTensors."
+            )
+
+        with set_ctx_getter(error_on_ctx):
+            return fake_impl_holder.kernel(*args, **kwargs)
+
+    return meta_kernel
+
+
+def get_none():
+    return None
+
+
+global_ctx_getter: Callable = get_none
+
+
+@contextlib.contextmanager
+def set_ctx_getter(ctx_getter):
+    global global_ctx_getter
+    prev = global_ctx_getter
+    try:
+        global_ctx_getter = ctx_getter
+        yield
+    finally:
+        global_ctx_getter = prev
+
+
+class FakeImplCtx:
+    """
+    Context object for writing fake implementations for custom operators.
+    """
+
+    def __init__(self, _fake_mode, _op):
+        self._fake_mode = _fake_mode
+        self._shape_env = _fake_mode.shape_env
+        self._op = _op
+
+    @deprecated(
+        "`create_unbacked_symint` is deprecated, please use `new_dynamic_size` instead",
+        category=FutureWarning,
+    )
+    def create_unbacked_symint(self, *, min=2, max=None) -> torch.SymInt:
+        return self.new_dynamic_size(min=min, max=max)
+
+    def new_dynamic_size(self, *, min=0, max=None) -> torch.SymInt:
+        """Constructs a new symint (symbolic int) representing a data-dependent value.
+
+        This is useful for writing the fake implementation (which is necessary
+        for torch.compile) for a CustomOp where an output Tensor has a size
+        that depends on the data of the input Tensors.
+
+        Args:
+            min (int): A statically known inclusive lower bound for this symint. Default: 0
+            max (Optional[int]): A statically known inclusive upper bound for this
+                symint. Default: None
+
+        .. warning:
+
+            It is important that the ``min`` and ``max`` (if not None) values are set
+            correctly, otherwise, there will be undefined behavior under
+            torch.compile. The default value of ``min`` is 2 due to torch.compile
+            specializing on 0/1 sizes.
+
+            You must also verify that your implementation on concrete Tensors
+            (e.g. CPU/CUDA) only returns Tensors where the size that corresponds
+            to the symint also has respects these constraint.
+            The easiest way to do this is to add an assertion in the CPU/CUDA/etc
+            implementation that the size follows these bounds.
+
+        Example::
+
+            >>> # An operator with data-dependent output shape
+            >>> lib = torch.library.Library("mymodule", "FRAGMENT")
+            >>> lib.define("mymodule::custom_nonzero(Tensor x) -> Tensor")
+            >>>
+            >>> @torch.library.register_fake("mymodule::custom_nonzero")
+            >>> def _(x):
+            >>>     # Number of nonzero-elements is data-dependent.
+            >>>     # Since we cannot peek at the data in an fake impl,
+            >>>     # we use the ctx object to construct a new symint that
+            >>>     # represents the data-dependent size.
+            >>>     ctx = torch.library.get_ctx()
+            >>>     nnz = ctx.new_dynamic_size()
+            >>>     shape = [nnz, x.dim()]
+            >>>     result = x.new_empty(shape, dtype=torch.int64)
+            >>>     return result
+            >>>
+            >>> @torch.library.impl(lib, "custom_nonzero", "CPU")
+            >>> def _(x):
+            >>>     x_np = x.numpy()
+            >>>     res = np.stack(np.nonzero(x_np), axis=1)
+            >>>     return torch.tensor(res, device=x.device)
+
+        """
+        if (
+            self._shape_env is None
+            or not self._shape_env.allow_dynamic_output_shape_ops
+        ):
+            raise torch._subclasses.fake_tensor.DynamicOutputShapeException(self._op)
+
+        if isinstance(min, torch.SymInt) or isinstance(max, torch.SymInt):
+            raise ValueError(
+                f"ctx.new_dynamic_size(min={min}, max={max}): expected "
+                f"min and max to be statically known ints but got SymInt. "
+                f"This is not supported."
+            )
+
+        if min < 0:
+            raise ValueError(
+                f"ctx.new_dynamic_size(min={min}, ...): expected min to be "
+                f"greater than or equal to 0: this API can only create "
+                f"non-negative sizes."
+            )
+
+        return allocate_size(self._shape_env, min, max)
+
+
+def allocate_size(shape_env, min_val=0, max_val=None):
+    result = shape_env.create_unbacked_symint()
+    torch.fx.experimental.symbolic_shapes._constrain_range_for_size(
+        result, min=min_val, max=max_val
+    )
+    return result

kronos-venv/lib/python3.12/site-packages/torch/_library/fake_profile.py

@@ -0,0 +1,327 @@
+import contextlib
+import io
+import logging
+import os
+from collections.abc import Callable, Generator
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+from torch._library.custom_ops import _maybe_get_opdef
+from torch.types import FileLike
+
+
+log = logging.getLogger(__name__)
+
+
+class MissingOpProfile(RuntimeError):
+    """
+    This is raised when we don't have an operator profile available for the
+    given inputs.
+    """
+
+
+@dataclass(frozen=True)
+class TensorMetadata:
+    rank: int
+    dtype: torch.dtype
+    device: torch.device
+    layout: torch.layout
+
+    @staticmethod
+    def maybe_from_tensor(t: Any) -> Optional["TensorMetadata"]:
+        if not isinstance(t, torch.Tensor):
+            return None
+        return TensorMetadata(t.dim(), t.dtype, t.device, t.layout)
+
+
+@dataclass(frozen=True)
+class OpProfile:
+    args_profile: tuple[TensorMetadata | None]
+    out_profile: TensorMetadata | tuple[TensorMetadata]
+
+
+def _generate_fake_kernel(op_name: str, op_profile: set[OpProfile]) -> Callable:
+    def _match_args(args_profile: tuple[TensorMetadata | None], args: Any) -> bool:
+        return all(
+            TensorMetadata.maybe_from_tensor(arg) == args_profile[i]
+            for i, arg in enumerate(args)
+        )
+
+    def _generate_res(
+        out_profile: TensorMetadata | tuple[TensorMetadata],
+    ) -> torch.Tensor | list[torch.Tensor]:
+        ctx = torch.library.get_ctx()
+
+        def _generate_tensor_out(t: TensorMetadata) -> torch.Tensor:
+            fake_shape = [ctx.new_dynamic_size() for _ in range(t.rank)]
+            fake_strides = [-1] * t.rank
+            expected = 1
+            fake_stride = expected
+            # pyrefly: ignore [bad-assignment]
+            for i in range(t.rank):
+                fake_strides[i] = fake_stride  # type: ignore[assignment]
+                fake_stride = fake_stride * fake_shape[i]  # type: ignore[assignment]
+
+            return torch.empty_strided(
+                fake_shape,
+                fake_strides,
+                device=t.device,
+                dtype=t.dtype,
+                layout=t.layout,
+            )
+
+        if isinstance(out_profile, TensorMetadata):
+            return _generate_tensor_out(out_profile)
+        else:
+            return [_generate_tensor_out(t) for t in out_profile]
+
+    def _fake_kernel(*args, **kwargs):  # type: ignore[no-untyped-def]
+        for profile in op_profile:
+            if _match_args(profile.args_profile, (*args, *kwargs.values())):
+                return _generate_res(profile.out_profile)
+
+        raise MissingOpProfile(
+            f"No fake kernel was found for {op_name}, and although we have "
+            "previously registered some profiles to generate a fake kernel, "
+            f"no profiles match the given inputs: {args, kwargs}."
+        )
+
+    return _fake_kernel
+
+
+@contextlib.contextmanager
+def unsafe_generate_fake_kernels(op_profiles: dict[str, set[OpProfile]]) -> Generator:
+    """
+    Registers a fake kernel based on the given operator profiles. This fake
+    kernel registration will override any existing fake kernel registrations.
+
+    The input is a dictionary mapping operator names to a set of operator
+    profiles, which we will use to generate fake kernels. The operator profiles
+    are a record of the input and output tensor metadata. Based on this
+    information we will match a given input to the recorded profile, and return
+    an output with the same metadata as in the recorded profile. If a profile
+    doesn't exist then an exception will be thrown.
+
+    The fake kernel generation is considered unsafe because it relies on the
+    rigid, pre-defined operator profiles that do not account for potential
+    variations in output behavior. Specifically, the generated kernels assume a
+    fixed relationship between input and output ranks. However, in reality, it's
+    possible that data-dependent operations may produce outputs of different
+    ranks even when given inputs of the same rank. The generated fake kernels
+    are inflexible and unable to accommodate these nuances, making them
+    potentially unsafe.
+
+    Args:
+        op_profiles (dict[str, set[OpProfile]]): A dictionary mapping operator
+            name to a set of operator profiles from which we will generate fake
+            kernels.
+
+    Examples:
+
+        >>> # Example: Registering an op-profile from draft-export
+        >>> import torch
+        >>> from torch.export._draft_export import draft_export
+        >>>
+        >>> @torch.library.custom_op("mylib::foo", mutates_args=())
+        >>> def foo(x: Tensor, y: Tensor) -> Tensor:
+        >>>     return x + y
+        >>>
+        >>> class M(torch.nn.Module):
+        >>>     def forward(self, a, b):
+        >>>         res = torch.ops.mylib.foo(a, b)  # no fake impl
+        >>>         return res
+        >>>
+        >>> ep = draft_export(M(), (torch.ones(3, 4), torch.ones(3, 4))
+        >>>
+        >>> with torch._library.fake_profile.unsafe_generate_fake_kernels(ep._report.op_profiles):
+        >>>     decomp = ep.run_decompositions()
+
+    """
+
+    libs: list[torch.library.Library] = []
+    # Stores old fake impls from custom ops declared through @custom_op
+    old_fake_impls: dict[str, Callable] = {}
+    for op_name, profiles in op_profiles.items():
+        log.warning(
+            "Registering fake profile for %s. This will override any existing "
+            "fake kernel registration.",
+            op_name,
+        )
+
+        op_name_split = op_name.split(".")
+        namespace, op_name_str = op_name_split[0], op_name_split[1]
+        op_str = f"{namespace}::{op_name_str}"
+
+        fake_kernel = _generate_fake_kernel(op_str, profiles)
+
+        if opdef := _maybe_get_opdef(op_str):
+            # If the op is a CustomOpDef, save the existing abstract_fn so that
+            # we can restore it after this contextmanager
+            if opdef._abstract_fn is not None:
+                old_fake_impls[op_str] = opdef._abstract_fn
+            opdef.register_fake(fake_kernel)
+
+        else:
+            # Create a new library so that we can register a new fake impl.
+            # These libraries will then be destroyed after the contextmanager,
+            # which will automatically restore the previously registered fake
+            # impls.
+            newlib = torch.library.Library(namespace, "FRAGMENT")  # noqa: TOR901
+            torch.library.register_fake(
+                op_str, fake_kernel, lib=newlib, allow_override=True
+            )
+            libs.append(newlib)
+
+    try:
+        yield libs
+    finally:
+        # Destroying the libraries will automatically restore the previously
+        # registered fake impls
+        for lib in libs:
+            lib._destroy()
+
+        # Restore abstract_fns for CustomOpDefs
+        for op_str, old_fake in old_fake_impls.items():
+            opdef = _maybe_get_opdef(op_str)
+            if opdef is None:
+                raise AssertionError(f"opdef for {op_str} must not be None")
+            opdef.register_fake(old_fake)
+
+
+def get_torch_version() -> str:
+    version = torch.__version__.split(".")
+    return f"{int(version[0])}.{int(version[1])}"
+
+
+def generate_yaml_from_profiles(op_profiles: dict[str, set[OpProfile]]) -> str:
+    """
+    Generates a yaml string from the given operator profiles which can be saved
+    to a file. The yaml string can be loaded back into an operator profile
+    structure using `read_profiles_from_yaml`.
+    """
+
+    import yaml
+
+    from torch._export.serde.serialize import (
+        _TORCH_TO_SERIALIZE_DTYPE,
+        _TORCH_TO_SERIALIZE_LAYOUT,
+    )
+
+    def serialize_tensor_metadata(t: TensorMetadata) -> dict:
+        return {
+            "rank": t.rank,
+            "dtype": _TORCH_TO_SERIALIZE_DTYPE[t.dtype].value,
+            "device": str(t.device),
+            "layout": _TORCH_TO_SERIALIZE_LAYOUT[t.layout].value,
+        }
+
+    def serialize_op_profile(op: OpProfile) -> dict:
+        return {
+            "args_profile": [
+                serialize_tensor_metadata(arg)
+                for arg in op.args_profile
+                if arg is not None
+            ],
+            "out_profile": (
+                serialize_tensor_metadata(op.out_profile)
+                if isinstance(op.out_profile, TensorMetadata)
+                else [serialize_tensor_metadata(out) for out in op.out_profile]
+            ),
+        }
+
+    serialized_data = {
+        operator: [serialize_op_profile(profile) for profile in profiles]
+        for operator, profiles in op_profiles.items()
+    }
+    return yaml.dump(
+        {"torch_version": get_torch_version(), "operators": serialized_data},
+        sort_keys=False,
+    )
+
+
+def save_op_profiles(op_profiles: dict[str, set[OpProfile]], f: FileLike) -> None:
+    """
+    Serializes the given operator profiles into a yaml format and saves it to
+    the given file. The operator profile can be loaded back using `load_op_profiles`.
+    """
+    yaml_str = generate_yaml_from_profiles(op_profiles)
+
+    if isinstance(f, (str, os.PathLike)):
+        f = os.fspath(f)
+
+        with open(f, "w") as file:
+            file.write(yaml_str)
+
+    elif isinstance(f, io.BytesIO):
+        f.write(yaml_str.encode("utf-8"))
+
+    else:
+        raise ValueError(f"Invalid type of file {f}")
+
+
+def read_profiles_from_yaml(yaml_str: str) -> dict[str, set[OpProfile]]:
+    """
+    Reads the yaml saved by `save_op_profiles` and returns the operator profiles.
+    """
+
+    import yaml
+
+    from torch._export.serde.serialize import (
+        _SERIALIZE_TO_TORCH_DTYPE,
+        _SERIALIZE_TO_TORCH_LAYOUT,
+    )
+
+    def deserialize_tensor_metadata(data: dict) -> TensorMetadata:
+        return TensorMetadata(
+            rank=data["rank"],
+            dtype=_SERIALIZE_TO_TORCH_DTYPE[data["dtype"]],
+            device=torch.device(data["device"]),
+            layout=_SERIALIZE_TO_TORCH_LAYOUT[data["layout"]],
+        )
+
+    def deserialize_op_profile(data: dict) -> OpProfile:
+        args_profile = tuple(
+            deserialize_tensor_metadata(arg) for arg in data["args_profile"]
+        )
+        out_profile_data = data["out_profile"]
+        out_profile: tuple[TensorMetadata] | TensorMetadata = (
+            tuple(deserialize_tensor_metadata(out) for out in out_profile_data)  # type: ignore[assignment]
+            if isinstance(out_profile_data, list)
+            else deserialize_tensor_metadata(out_profile_data)
+        )
+        return OpProfile(args_profile=args_profile, out_profile=out_profile)  # type: ignore[arg-type]
+
+    loaded_data = yaml.safe_load(yaml_str)
+    loaded_torch_version = loaded_data["torch_version"]
+
+    if loaded_torch_version != get_torch_version():
+        raise RuntimeError(
+            "Unable to load outdated profile. It was saved with torch version: "
+            f"{loaded_torch_version} but the current torch version is: {get_torch_version()}"
+        )
+
+    operators_data = loaded_data["operators"]
+    return {
+        operator: {deserialize_op_profile(profile) for profile in profiles}
+        for operator, profiles in operators_data.items()
+    }
+
+
+def load_op_profiles(f: FileLike) -> dict[str, set[OpProfile]]:
+    """
+    Loads the saved operator profiles from `save_op_profiles`.
+    """
+    if isinstance(f, (str, os.PathLike)):
+        f = os.fspath(f)
+
+        with open(f) as file:
+            yaml_str = file.read()
+
+    elif isinstance(f, io.BytesIO):
+        yaml_str = f.read().decode("utf-8")
+
+    else:
+        raise ValueError(f"Invalid type of file {f}")
+
+    return read_profiles_from_yaml(yaml_str)

kronos-venv/lib/python3.12/site-packages/torch/_library/infer_schema.py

@@ -0,0 +1,367 @@
+# mypy: allow-untyped-defs
+import collections
+import inspect
+import typing
+from types import GenericAlias
+
+import torch
+from torch import device, dtype, Tensor, types
+from torch.utils._exposed_in import exposed_in
+
+from .opaque_object import (
+    _resolve_opaque_type_info,
+    is_opaque_reference_type,
+    is_opaque_type,
+)
+
+
+# This is used as a negative test for
+# test_custom_ops.py::TestTypeConversion::test_type_eval.
+_TestTensor = torch.Tensor
+
+
+@exposed_in("torch.library")
+def infer_schema(
+    prototype_function: typing.Callable,
+    /,
+    *,
+    mutates_args,
+    op_name: str | None = None,
+) -> str:
+    r"""Parses the schema of a given function with type hints. The schema is inferred from the
+    function's type hints, and can be used to define a new operator.
+
+    We make the following assumptions:
+
+    * None of the outputs alias any of the inputs or each other.
+    * | String type annotations "device, dtype, Tensor, types" without library specification are
+      | assumed to be torch.*. Similarly, string type annotations "Optional, List, Sequence, Union"
+      | without library specification are assumed to be typing.*.
+    * | Only the args listed in ``mutates_args`` are being mutated. If ``mutates_args`` is "unknown",
+      | it assumes that all inputs to the operator are being mutates.
+
+    Callers (e.g. the custom ops API) are responsible for checking these assumptions.
+
+    Args:
+        prototype_function: The function from which to infer a schema for from its type annotations.
+        op_name (Optional[str]): The name of the operator in the schema. If ``name`` is None, then the
+            name is not included in the inferred schema. Note that the input schema to
+            ``torch.library.Library.define`` requires a operator name.
+        mutates_args ("unknown" | Iterable[str]): The arguments that are mutated in the function.
+
+    Returns:
+        The inferred schema.
+
+    Example:
+        >>> def foo_impl(x: torch.Tensor) -> torch.Tensor:
+        >>>     return x.sin()
+        >>>
+        >>> infer_schema(foo_impl, op_name="foo", mutates_args={})
+        foo(Tensor x) -> Tensor
+        >>>
+        >>> infer_schema(foo_impl, mutates_args={})
+        (Tensor x) -> Tensor
+    """
+    UNKNOWN_MUTATES = "unknown"
+    pf_globals = prototype_function.__globals__
+    pf_locals = None
+    # TODO: Once our minimum version is py3.10+ pass `eval_str=True` to
+    # inspect.signature() and we no longer need to deal with stringified
+    # annotations below.
+    sig = inspect.signature(prototype_function)
+
+    def error_fn(what):
+        raise ValueError(f"infer_schema(func): {what} Got func with signature {sig})")
+
+    def convert_type_string(annotation_type: str):
+        try:
+            return eval(annotation_type, pf_globals, pf_locals)
+        except Exception:
+            error_fn(
+                f"Unsupported type annotation {annotation_type}. It is not a type."
+            )
+
+    def unstringify_types(
+        tys: tuple[type[object] | str, ...],
+    ) -> tuple[tuple[typing.Any, ...], bool]:
+        res = []
+        changed = False
+        for ty in tys:
+            ty, ty_changed = unstringify_type(ty)
+            res.append(ty)
+            changed |= ty_changed
+        if changed:
+            return tuple(res), True
+        else:
+            return tys, False  # type: ignore[return-value]
+
+    def unstringify_type(ty: type[object] | str) -> tuple[typing.Any, bool]:
+        # Dig through a generic type and if it contains a stringified type
+        # convert that to a real type. The second return value indicates if the
+        # type contained a string or not.
+        if isinstance(ty, str):
+            return convert_type_string(ty), True
+        elif origin := typing.get_origin(ty):
+            args, args_changed = unstringify_types(typing.get_args(ty))
+            if args_changed:
+                return GenericAlias(origin, args), True
+
+        return ty, False
+
+    params = []
+    seen_args = set()
+    saw_kwarg_only_arg = False
+    for idx, (name, param) in enumerate(sig.parameters.items()):
+        if not supported_param(param):
+            error_fn("We do not support positional-only args, varargs, or varkwargs.")
+
+        if param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # The first time we see a kwarg-only arg, add "*" to the schema.
+            if not saw_kwarg_only_arg:
+                params.append("*")
+                saw_kwarg_only_arg = True
+
+        if param.annotation is inspect.Parameter.empty:
+            error_fn(f"Parameter {name} must have a type annotation.")
+
+        # The annotation might be converted to a string by annotation,
+        # we convert it to the actual type.
+        annotation_type, _ = unstringify_type(param.annotation)
+
+        schema_type = None
+        if annotation_type not in SUPPORTED_PARAM_TYPES:
+            if is_opaque_type(annotation_type):
+                schema_type = _resolve_opaque_type_info(annotation_type).class_name  # type: ignore[union-attr]
+            elif annotation_type == torch._C.ScriptObject:
+                error_fn(
+                    f"Parameter {name}'s type cannot be inferred from the schema "
+                    "as it is a ScriptObject. Please manually specify the schema "
+                    "using the `schema=` kwarg with the actual type of the ScriptObject."
+                )
+            elif (
+                hasattr(annotation_type, "__origin__")
+                and annotation_type.__origin__ is tuple
+            ):
+                list_type = tuple_to_list(annotation_type)
+                example_type_str = "\n\n"
+                # Only suggest the list type if this type is supported.
+                if list_type in SUPPORTED_PARAM_TYPES:
+                    example_type_str = f"For example, {list_type}.\n\n"
+                error_fn(
+                    f"Parameter {name} has unsupported type {param.annotation}. "
+                    f"We do not support Tuple inputs in schema. As a workaround, please try to use List instead. "
+                    f"{example_type_str}"
+                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
+                )
+            else:
+                error_fn(
+                    f"Parameter {name} has unsupported type {param.annotation}. "
+                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
+                )
+        else:
+            schema_type = SUPPORTED_PARAM_TYPES[annotation_type]
+
+        if schema_type is None:
+            raise AssertionError(f"schema_type is None for param {name}")
+
+        if type(mutates_args) is str:
+            if mutates_args != UNKNOWN_MUTATES:
+                raise ValueError(
+                    "mutates_args must either be a sequence of the names of "
+                    "the arguments that are mutated or the string 'unknown'. "
+                )
+            if schema_type.startswith("Tensor"):
+                schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor') :]}"
+        elif name in mutates_args:
+            if not schema_type.startswith("Tensor"):
+                error_fn(
+                    f"Parameter {name} is in mutable_args but only Tensors or collections of Tensors can be mutated"
+                )
+            schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor') :]}"
+        seen_args.add(name)
+        if param.default is inspect.Parameter.empty:
+            # pyrefly: ignore [bad-argument-type]
+            params.append(f"{schema_type} {name}")
+        else:
+            default_repr = None
+            if param.default is None or isinstance(param.default, (int, float, bool)):
+                default_repr = str(param.default)
+            elif isinstance(param.default, (str, torch.device)):
+                default_repr = f'"{param.default}"'
+            elif isinstance(param.default, torch.dtype):
+                dtype_repr = str(param.default)
+                torch_dot = "torch."
+                if not dtype_repr.startswith(torch_dot):
+                    raise AssertionError(
+                        f"dtype repr {dtype_repr!r} must start with 'torch.'"
+                    )
+                default_repr = dtype_repr[len(torch_dot) :]
+            else:
+                error_fn(
+                    f"Parameter {name} has an unsupported default value type {type(param.default)}. "
+                    f"Please file an issue on GitHub so we can prioritize this."
+                )
+            # pyrefly: ignore [bad-argument-type]
+            params.append(f"{schema_type} {name}={default_repr}")
+    if mutates_args != UNKNOWN_MUTATES:
+        mutates_args_not_seen = set(mutates_args) - seen_args
+        if len(mutates_args_not_seen) > 0:
+            error_fn(
+                f"{mutates_args_not_seen} in mutates_args were not found in "
+                f"the custom op's signature. "
+                f"mutates_args should contain the names of all args that the "
+                f"custom op mutates, or just the string 'unknown' if you don't know."
+            )
+    return_annotation, _ = unstringify_type(sig.return_annotation)
+    ret = parse_return(return_annotation, error_fn)
+    if op_name is not None:
+        return f"{op_name}({', '.join(params)}) -> {ret}"
+    return f"({', '.join(params)}) -> {ret}"
+
+
+def derived_types(
+    base_type: type | typing._SpecialForm,
+    cpp_type: str,
+    list_base: bool,
+    optional_base_list: bool,
+    optional_list_base: bool,
+):
+    result: list[tuple[type | typing._SpecialForm | GenericAlias, str]] = [
+        (base_type, cpp_type),
+        # pyrefly: ignore [not-a-type]
+        (typing.Optional[base_type], f"{cpp_type}?"),  # noqa: UP045
+    ]
+
+    def derived_seq_types(typ: type | typing._SpecialForm):
+        return (
+            typing.Sequence[typ],  # type: ignore[valid-type]  # noqa: UP006
+            typing.List[typ],  # type: ignore[valid-type]  # noqa: UP006
+            GenericAlias(collections.abc.Sequence, (typ,)),
+            GenericAlias(list, (typ,)),
+        )
+
+    if list_base:
+        result.extend(
+            (seq_typ, f"{cpp_type}[]") for seq_typ in derived_seq_types(base_type)
+        )
+    if optional_base_list:
+        result.extend(
+            (seq_typ, f"{cpp_type}?[]")
+            # pyrefly: ignore [not-a-type]
+            for seq_typ in derived_seq_types(typing.Optional[base_type])  # noqa: UP045
+        )
+    if optional_list_base:
+        result.extend(
+            (typing.Optional[seq_typ], f"{cpp_type}[]?")  # noqa: UP045
+            for seq_typ in derived_seq_types(base_type)
+        )
+    return result
+
+
+def get_supported_param_types():
+    data: list[tuple[type | typing._SpecialForm, str, bool, bool, bool]] = [
+        # (python type, schema type, type[] variant, type?[] variant, type[]? variant
+        (Tensor, "Tensor", True, True, False),
+        (int, "SymInt", True, False, True),
+        (float, "float", True, False, True),
+        (bool, "bool", True, False, True),
+        (str, "str", False, False, False),
+        (types.Number, "Scalar", True, False, False),
+        (dtype, "ScalarType", False, False, False),
+        (device, "Device", False, False, False),
+    ]
+
+    if torch.distributed.is_available():
+        from torch.distributed.distributed_c10d import GroupName
+
+        data.append((typing.cast(type, GroupName), "str", False, False, False))
+
+    result = []
+    for line in data:
+        result.extend(derived_types(*line))
+    return dict(result)
+
+
+SUPPORTED_RETURN_TYPES = {
+    Tensor: "Tensor",
+    typing.List[Tensor]: "Tensor[]",  # noqa: UP006
+    list[Tensor]: "Tensor[]",
+    int: "SymInt",
+    float: "float",
+    bool: "bool",
+    types.Number: "Scalar",
+}
+
+
+def parse_return(annotation, error_fn):
+    if annotation is None:
+        return "()"
+
+    if annotation is inspect.Parameter.empty:
+        error_fn("No return type annotation was provided. Please add one.")
+
+    origin = typing.get_origin(annotation)
+    if origin is not tuple:
+        if annotation not in SUPPORTED_RETURN_TYPES:
+            if is_opaque_reference_type(annotation):
+                return _resolve_opaque_type_info(annotation).class_name  # type: ignore[union-attr]
+            error_fn(
+                f"Return has unsupported type {annotation}. "
+                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
+            )
+
+        return SUPPORTED_RETURN_TYPES[annotation]
+
+    args = typing.get_args(annotation)
+    for arg in args:
+        if arg not in SUPPORTED_RETURN_TYPES and not is_opaque_reference_type(arg):
+            error_fn(
+                f"Return has unsupported type {annotation}. "
+                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
+            )
+
+    def _return_type_str(arg):
+        if ty := SUPPORTED_RETURN_TYPES.get(arg):
+            return ty
+        return _resolve_opaque_type_info(arg).class_name  # type: ignore[union-attr]
+
+    output_ty = ", ".join(_return_type_str(arg) for arg in args)
+
+    # use (()) to represent tuple with single element
+    if len(args) == 1:
+        output_ty = "(" + output_ty + ")"
+    return "(" + output_ty + ")"
+
+
+SUPPORTED_PARAM_TYPES = get_supported_param_types()
+
+
+def supported_param(param: inspect.Parameter) -> bool:
+    return param.kind in (
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.KEYWORD_ONLY,
+    )
+
+
+def tuple_to_list(tuple_type: type[tuple]) -> type[list]:
+    """
+    Convert `tuple_type` into a list type with the same type arguments. Assumes that `tuple_type` is typing.Tuple type.
+    """
+    type_args = getattr(tuple_type, "__args__", None)
+    # Account for different python versions, e.g. python 3.8 would give ()
+    # but python 3.12 would give None.
+    if (
+        tuple_type is typing.Tuple  # noqa: UP006
+        or tuple_type is tuple
+        or type_args == ()
+        or type_args is None
+    ):
+        # Handle the case of an empty tuple type
+        return list
+    elif len(type_args) == 1:
+        # General case: create a List with the same type arguments
+        return list[type_args[0]]  # type: ignore[valid-type]
+    elif len(type_args) == 2 and type_args[1] is Ellipsis:
+        return list[type_args[0]]  # type: ignore[valid-type]
+    else:
+        return list[typing.Union[tuple(type_args)]]  # type: ignore[misc, return-value]  # noqa: UP007

kronos-venv/lib/python3.12/site-packages/torch/_library/opaque_object.py

@@ -0,0 +1,422 @@
+"""
+Note [Opaque Objects]
+
+Opaque objects are the way we allow custom operators to accept a user-defined
+"black box" object as an input.
+
+There are two kinds of opaque types: VALUE type and REFERENCE type.
+The distinction determines how torch.compile handles the object.
+
+REFERENCE TYPES (default):
+
+Reference-typed opaque objects represent mutable stateful objects and are
+treated as black boxes. In torch.compile, since torch.compile cannot optimize
+the anything (including tensors) within the object, the object must be an
+input to the graph.
+
+You can register a custom class as being a reference-based opaque object class
+through `register_opaque_type(MyClass, typ="reference")`.
+
+VALUE TYPES:
+
+Value-typed opaque objects represent constant values.
+In torch.compile, the graph specializes on the object like how other constants
+are. Therefore there are a couple of methods on the class that must be
+implemented before registering it as a value-typed opaque object class:
+  - __eq__: torch.compile will create guards based on the equality of this
+  object, meaning that a recompilation will happen if __eq__ returns False.
+  - __hash__: This must be implemented for Fake Tensor caching
+  - __fx_repr__: This must be implemented to provide an evaluable representation
+    for FX graph codegen. It should return a tuple of (repr_string, dict[str, type])
+    where repr_string can reconstruct the object and the dict maps names used in
+    repr_string to their corresponding types.
+
+You can register a custom class as being a reference-based opaque object class
+through `register_opaque_type(MyClass, typ="value")`.
+"""
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Literal, NewType, TYPE_CHECKING, TypeAlias
+from typing_extensions import TypeIs
+from weakref import WeakKeyDictionary
+
+import torch
+from torch._opaque_base import OpaqueBase, OpaqueBaseMeta  # noqa: F401
+
+
+if TYPE_CHECKING:
+    from torch.fx import Proxy
+    from torch.fx.experimental.proxy_tensor import PythonKeyTracer
+
+from .fake_class_registry import register_fake_class
+
+
+log = logging.getLogger(__name__)
+
+
+class MemberType(Enum):
+    """
+    Defines how a member (attribute/property/method) of an opaque object is handled
+    during torch.compile tracing.
+    """
+
+    # Reads/calls the member at trace time with the real object and bakes the result as a constant
+    USE_REAL = "use_real"
+    # Inlines/traces the member
+    INLINED = "inlined"
+
+
+@register_fake_class("aten::OpaqueObject")
+class FakeOpaqueObject:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def __obj_unflatten__(cls, flattened_ctx: dict[str, Any]) -> None:
+        raise RuntimeError(
+            "FakeOpaqueObject should not be created through __obj_unflatten__ "
+            "and should be special handled. Please file an issue to Github."
+        )
+
+
+OpaqueTypeStr = "__torch__.torch.classes.aten.OpaqueObject"
+
+OpaqueType = NewType("OpaqueType", torch._C.ScriptObject)
+
+# Type for reconstruct_fn: called by PythonKeyTracer.create_arg when make_fx
+# encounters an untracked opaque reference (e.g. a backward closure capture).
+# Should derive the object from existing graph inputs or return None to fall
+# back to get_attr.  Args: (obj, get_tracked_proxy, tracer).
+ReconstructFn: TypeAlias = Callable[
+    [OpaqueBase, Callable[[OpaqueBase], "Proxy | None"], "PythonKeyTracer"],
+    "Proxy | None",
+]
+
+
+@dataclass
+class _OpaqueTypeInfo:
+    class_name: str
+    opaque_typ: Literal["reference", "value"]
+    guard_fn: Callable[
+        [Any], list[Any]
+    ]  # Callable that takes the object and returns list of values to guard on
+    members: dict[str, MemberType]  # Maps member name to how it should be handled
+    hoist: bool
+    reconstruct_fn: ReconstructFn | None
+
+
+# Mapping of type -> (string name, reference/value type)
+_OPAQUE_TYPES: WeakKeyDictionary[Any, _OpaqueTypeInfo] = WeakKeyDictionary()
+# Mapping of class_name -> (type, reference/value type)
+_OPAQUE_TYPES_BY_NAME: dict[str, _OpaqueTypeInfo] = {}
+
+
+def _resolve_opaque_type_info(cls: Any) -> _OpaqueTypeInfo | None:
+    if cls in _OPAQUE_TYPES:
+        return _OPAQUE_TYPES[cls]
+    if not isinstance(cls, type):
+        return None
+
+    # Allow subclasses too
+    for parent in cls.__mro__[1:]:
+        if parent in _OPAQUE_TYPES:
+            return _OPAQUE_TYPES[parent]
+    return None
+
+
+def get_opaque_type_name(cls: Any) -> str:
+    """
+    Gets the registered opaque type name for a given class.
+
+    Args:
+        cls (type): The class to get the type name for.
+
+    Returns:
+        str: The registered type name for the class.
+
+    Raises:
+        ValueError: If the class is not registered as an opaque type.
+    """
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        raise ValueError(
+            f"Class {cls} is not registered as an opaque type. "
+            f"Call register_opaque_type({cls.__name__}) first."
+        )
+    return info.class_name
+
+
+def register_opaque_type(
+    cls: Any,
+    *,
+    typ: str,
+    hoist=False,
+    guard_fn: Any = None,
+    members: dict[str, MemberType] | None = None,
+    reconstruct_fn: ReconstructFn | None = None,
+) -> None:
+    """
+    Registers the given type as an opaque type which allows this to be consumed
+    by a custom operator.
+
+    The type name will be automatically generated from the class's fully
+    qualified name (ex. my_module.MyClass).
+
+    Args:
+        cls (type): The class to register as an opaque type.
+        typ (str): Either "reference" or "value". See Note [Opaque Objects] for
+            more details.
+        hoist (bool): Only applies to value types. A hoist=True value type
+            object is lifted as an input to the torch.compile'd graph, instead
+            of being a constant baked into the graph. This is useful to
+            improve compilation times in hierarchical compilation
+            (e.g., change your custom ops to use hoisted strings to avoid
+            baking the string into the Dynamo/AOTAutograd/FX graphs).
+            This flag does nothing for reference types.
+        guard_fn (callable | None): A function that takes an instance of the opaque
+            object and returns a list of values to guard on. These values will be compared
+            for equality on each function call, triggering recompilation if they change.
+            Only applicable for reference types.
+            Example: lambda obj: [obj.x, obj.y]
+        members (dict[str, MemberType] | None): Dictionary mapping member names
+            (attributes, properties, or methods) to their MemberType, which controls
+            how they are handled during torch.compile tracing:
+            - MemberType.USE_REAL: Evaluates with the real object at compile time and
+              bakes the result as a constant
+            - MemberType.INLINED: Inlines the method call into the trace
+    """
+    import torch.utils._pytree as pytree
+
+    # Prevent registration of built-in types (int, str, list, dict, etc.) and torch.Tensor
+    if cls.__module__ == "builtins" or cls is torch.Tensor:
+        raise ValueError(
+            f"Unable to register built-in type {cls} as an opaque type. "
+            "Please wrap it in a custom class and register the custom class as opaque."
+        )
+
+    if cls in pytree.SUPPORTED_NODES:
+        raise ValueError(
+            f"{cls} cannot be registered as an opaque object as it has been "
+            "registered as a pytree. Opaque objects must be pytree leaves."
+        )
+
+    # Value types store the real object directly during tracing (no
+    # FakeScriptObject wrapper), so they don't need OpaqueBaseMeta.
+    if typ != "value" and not isinstance(cls, OpaqueBaseMeta):
+        raise TypeError(
+            f"Opaque type {cls} must subclass torch._opaque_base.OpaqueBase "
+            "or 'metaclass=torch._opaque_base.OpaqueBaseMeta'. "
+            "This is required so that FakeScriptObject can be registered "
+            "as a virtual subclass, allowing isinstance() checks to work "
+            "during torch.compile tracing. "
+        )
+
+    if typ not in ["reference", "value"]:
+        raise AssertionError(
+            f"Opaque type must be either 'reference' or 'value', got {typ!r}"
+        )
+
+    if typ == "value":
+        # Enums use identity-based equality (singletons), which is fine for guarding.
+        if not issubclass(cls, Enum) and cls.__eq__ is object.__eq__:  # type: ignore[comparison-overlap]
+            raise TypeError(
+                f"Value-type opaque object of type {cls} is "
+                "expected to have a non-default `__eq__` "
+                "implementation as we will use this in torch.compile "
+                "to guard on the equality of objects."
+            )
+
+        # Class with a custom `__eq__` without `__hash__` won't inherit the default
+        # `__hash__` from object; see https://stackoverflow.com/a/1608907.
+        if cls.__hash__ is None:  # type: ignore[comparison-overlap]
+            raise TypeError(
+                f"Value-type opaque object of type {cls} is "
+                "expected to have a non-default `__hash__` "
+                "implementation as we will use this in torch.compile "
+                "for FakeTensor caching."
+            )
+
+        # Enums are special-cased in get_opaque_obj_repr.
+        if not issubclass(cls, Enum) and not hasattr(cls, "__fx_repr__"):
+            raise TypeError(
+                f"Value-type opaque object of type {cls} is "
+                "expected to have a `__fx_repr__` method "
+                "implementation as we will use this to reconstruct "
+                "the object in the FX codegen. __fx_repr__ should return "
+                "a tuple of (repr_string, dict[str, type])."
+            )
+
+        if guard_fn is not None:
+            raise TypeError(
+                "No need to specify `guard_fn` for "
+                f"value-type opaque class {cls} as it will be guarded based "
+                "on `__eq__`."
+            )
+
+    # Generate a fully qualified name by combining module and qualname
+    name = f"{cls.__module__}.{cls.__qualname__}"
+
+    type_info = _OpaqueTypeInfo(
+        name, typ, guard_fn, members or {}, hoist, reconstruct_fn
+    )
+    _OPAQUE_TYPES[cls] = type_info
+    _OPAQUE_TYPES_BY_NAME[name] = type_info
+
+    torch._C._register_opaque_type(name)
+
+
+# Enums are always opaque value types.
+register_opaque_type(Enum, typ="value")
+
+
+def is_opaque_value(value: object) -> TypeIs[OpaqueType]:
+    return is_opaque_type(type(value))
+
+
+def should_hoist(cls: Any) -> bool:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return info.hoist
+
+
+def get_reconstruct_fn(cls: type[OpaqueBase]) -> ReconstructFn | None:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return None
+    return info.reconstruct_fn
+
+
+def has_members(cls: Any) -> bool:
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return len(info.members) > 0
+
+
+def is_opaque_type(cls: type[Any] | str) -> bool:
+    """
+    Checks if the given type is an opaque type.
+    Also returns True for subclasses of registered opaque types.
+    """
+    if isinstance(cls, str):
+        return torch._C._is_opaque_type_registered(cls)
+
+    if not isinstance(cls, type):
+        log.warning("Passed invalid type `%s` to is_opaque_type, returning False", cls)
+        return False
+
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+
+    return torch._C._is_opaque_type_registered(info.class_name)
+
+
+def is_opaque_value_type(cls: type[Any] | str) -> bool:
+    """
+    Checks if the given type is an opaque **value** type.
+    See Note [Opaque Objects] for more information.
+    """
+    if not is_opaque_type(cls):
+        return False
+
+    if isinstance(cls, str):
+        return _OPAQUE_TYPES_BY_NAME[cls].opaque_typ == "value"
+
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return info.opaque_typ == "value"
+
+
+def is_opaque_reference_type(cls: Any) -> bool:
+    """
+    Checks if the given type is an opaque **reference** type.
+    See Note [Opaque Objects] for more information.
+    """
+    if not is_opaque_type(cls):
+        return False
+
+    if isinstance(cls, str):
+        return _OPAQUE_TYPES_BY_NAME[cls].opaque_typ == "reference"
+
+    info = _resolve_opaque_type_info(cls)
+    if info is None:
+        return False
+    return info.opaque_typ == "reference"
+
+
+def get_opaque_obj_repr(obj: Any) -> tuple[str, dict[str, type]]:
+    """
+    Get the FX-evaluable repr for an opaque object and collect required globals.
+
+    Objects must implement __fx_repr__() which should return:
+        (repr_string, dict_mapping_name_to_type)
+
+    where repr_string is an evaluable string representation and
+    dict_mapping_name_to_type maps the names used in repr_string to their types.
+
+    For example, if repr_string is "Foo(bar=Bar(1))", the dict should be:
+        {"Foo": Foo, "Bar": Bar}
+    """
+
+    # Enums are special cased
+    if isinstance(obj, Enum):
+        cls = type(obj)
+        return f"{cls.__name__}.{obj.name}", {cls.__name__: cls}
+
+    if not hasattr(obj, "__fx_repr__"):
+        raise TypeError(
+            f"Value-type opaque object of type {obj} is "
+            "expected to have a `__fx_repr__` method "
+            "implementation as we will use this to reconstruct "
+            "the object in the FX codegen. __fx_repr__ should return "
+            "a tuple of (repr_string, dict[str, type])."
+        )
+
+    repr_str, globals_dict = obj.__fx_repr__()
+
+    if not isinstance(repr_str, str):
+        raise TypeError(
+            f"__fx_repr__ for {type(obj).__name__} must return a string as the "
+            f"first element, got {type(repr_str).__name__}"
+        )
+
+    if not isinstance(globals_dict, dict):
+        raise TypeError(
+            f"__fx_repr__ for {type(obj).__name__} must return a dict as the "
+            f"second element, got {type(globals_dict).__name__}"
+        )
+
+    return repr_str, globals_dict
+
+
+def get_opaque_obj_info(cls: Any) -> _OpaqueTypeInfo | None:
+    if not is_opaque_type(cls):
+        return None
+
+    if isinstance(cls, str):
+        return _OPAQUE_TYPES_BY_NAME[cls]
+
+    return _resolve_opaque_type_info(cls)
+
+
+def get_member_type(cls: Any, member_name: str) -> MemberType | None:
+    """
+    Get the MemberType for a specific member of an opaque object class.
+
+    Args:
+        cls: The opaque object class (or its string name)
+        member_name: The name of the member to query
+
+    Returns:
+        MemberType if the member is registered, None otherwise
+    """
+    info = get_opaque_obj_info(cls)
+    if info is None:
+        return None
+    return info.members.get(member_name)

kronos-venv/lib/python3.12/site-packages/torch/_library/simple_registry.py

@@ -0,0 +1,91 @@
+from collections.abc import Callable
+from typing import Any
+
+from .effects import EffectHolder
+from .fake_impl import FakeImplHolder
+from .utils import RegistrationHandle
+
+
+__all__ = ["SimpleLibraryRegistry", "SimpleOperatorEntry", "singleton"]
+
+
+class SimpleLibraryRegistry:
+    """Registry for the "simple" torch.library APIs
+
+    The "simple" torch.library APIs are a higher-level API on top of the
+    raw PyTorch DispatchKey registration APIs that includes:
+    - fake impl
+
+    Registrations for these APIs do not go into the PyTorch dispatcher's
+    table because they may not directly involve a DispatchKey. For example,
+    the fake impl is a Python function that gets invoked by FakeTensor.
+    Instead, we manage them here.
+
+    SimpleLibraryRegistry is a mapping from a fully qualified operator name
+    (including the overload) to SimpleOperatorEntry.
+    """
+
+    def __init__(self) -> None:
+        self._data: dict[str, SimpleOperatorEntry] = {}
+
+    def find(self, qualname: str) -> "SimpleOperatorEntry":
+        res = self._data.get(qualname, None)
+        if res is None:
+            self._data[qualname] = res = SimpleOperatorEntry(qualname)
+        return res
+
+
+singleton: SimpleLibraryRegistry = SimpleLibraryRegistry()
+
+
+class SimpleOperatorEntry:
+    """This is 1:1 to an operator overload.
+
+    The fields of SimpleOperatorEntry are Holders where kernels can be
+    registered to.
+    """
+
+    def __init__(self, qualname: str) -> None:
+        self.qualname: str = qualname
+        self.fake_impl: FakeImplHolder = FakeImplHolder(qualname)
+        self.torch_dispatch_rules: GenericTorchDispatchRuleHolder = (
+            GenericTorchDispatchRuleHolder(qualname)
+        )
+
+        self.effect: EffectHolder = EffectHolder(qualname)
+
+    # For compatibility reasons. We can delete this soon.
+    @property
+    def abstract_impl(self) -> FakeImplHolder:
+        return self.fake_impl
+
+
+class GenericTorchDispatchRuleHolder:
+    def __init__(self, qualname: str) -> None:
+        self._data: dict[type, Callable[..., Any]] = {}
+        self.qualname: str = qualname
+
+    def register(
+        self, torch_dispatch_class: type, func: Callable[..., Any]
+    ) -> RegistrationHandle:
+        if self.find(torch_dispatch_class):
+            raise RuntimeError(
+                f"{torch_dispatch_class} already has a `__torch_dispatch__` rule registered for {self.qualname}"
+            )
+        self._data[torch_dispatch_class] = func
+
+        def deregister() -> None:
+            del self._data[torch_dispatch_class]
+
+        return RegistrationHandle(deregister)
+
+    def find(self, torch_dispatch_class: type) -> Callable[..., Any] | None:
+        return self._data.get(torch_dispatch_class, None)
+
+
+def find_torch_dispatch_rule(
+    op: Any, torch_dispatch_class: type
+) -> Callable[..., Any] | None:
+    return singleton.find(op.__qualname__).torch_dispatch_rules.find(
+        torch_dispatch_class
+    )

kronos-venv/lib/python3.12/site-packages/torch/_library/triton.py

@@ -0,0 +1,577 @@
+import ast
+import contextlib
+import inspect
+import logging
+import threading
+from collections.abc import Callable, Generator, Iterable
+from typing import Any
+
+from torch.utils._exposed_in import exposed_in
+
+from .custom_ops import custom_op, CustomOpDef
+from .infer_schema import infer_schema
+
+
+logger = logging.getLogger(__name__)
+
+triton_ops_to_kernels: dict[str, list[object]] = {}
+
+
+def get_triton_kernels_for_op(name: str) -> list[object]:
+    return triton_ops_to_kernels.get(name, [])
+
+
+def get_inner_triton_kernels(fn: Callable[..., Any]) -> list[object]:
+    """
+    Inspect the source of an arbitrary callable passed to torch._library.triton_op,
+    and grab all of the triton kernels that are wrapped inside of it.
+
+    This function traces local variable assignments to handle patterns like:
+        kernel_fn = _my_kernel  # global JITFunction
+        wrapped = some_wrapper(kernel_fn)
+        capture_triton(wrapped)[grid](...)
+
+    It also recursively analyzes called functions to find triton kernels hidden
+    behind helper function calls.
+
+    That said, it is best effort. There are cases (e.g., recursion > MAX_RECURSION_DEPTH)
+    that are not accounted for, so keep that in mind.
+    """
+
+    # prevent infinite recursion
+    MAX_RECURSION_DEPTH = 5
+
+    def find_triton_kernels(
+        fn: Callable[..., Any],
+        visited_fns: set[int] | None = None,
+        depth: int = 0,
+    ) -> list[object]:
+        try:
+            from triton.runtime.autotuner import Autotuner
+            from triton.runtime.jit import JITFunction
+        except ImportError:
+            logger.warning("Triton not available, find_triton_kernels = []")
+            return []
+
+        # unwrap decorated fn's (e.g., @lru_cache) to get the original
+        fn = inspect.unwrap(fn)
+
+        # init visited set and check for cycles/depth limit
+        if visited_fns is None:
+            visited_fns = set()
+
+        fn_id = id(fn)
+        if fn_id in visited_fns:
+            return []
+        if depth > MAX_RECURSION_DEPTH:
+            logger.debug(
+                "reached max recursion depth (%s) in find_triton_kernels",
+                MAX_RECURSION_DEPTH,
+            )
+            return []
+
+        visited_fns.add(fn_id)
+
+        try:
+            source = inspect.getsource(fn)
+        except (OSError, TypeError):
+            return []  # Source code not available
+
+        from torch._inductor.utils import IndentedBuffer
+
+        buffer = IndentedBuffer()
+        buffer.splice(source, strip=True)
+        tree = ast.parse(buffer.getrawvalue())
+
+        # Visitor to collect function calls, assignments, and triton kernels
+        class Visitor(ast.NodeVisitor):
+            def __init__(self) -> None:
+                self.triton_kernels: list[Any] = []
+                # track local variable assignments: var_name -> list of RHS expressions
+                self.assignments: dict[str, list[ast.expr]] = {}
+                # track function calls
+                self.called_functions: list[str] = []
+                # track return statement expressions
+                self.return_exprs: list[ast.expr] = []
+
+            def visit_Assign(self, node: ast.Assign) -> None:
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        self.assignments.setdefault(target.id, []).append(node.value)
+                self.generic_visit(node)
+
+            def visit_Return(self, node: ast.Return) -> None:
+                if node.value is not None:
+                    self.return_exprs.append(node.value)
+                self.generic_visit(node)
+
+            def visit_Call(self, node: ast.Call) -> None:
+                triton_func_names = ("capture_triton", "wrap_triton")
+                if isinstance(node.func, ast.Attribute):
+                    attr = node.func
+                    if isinstance(attr.value, ast.Attribute):
+                        if (
+                            isinstance(attr.value.value, ast.Name)
+                            and attr.value.value.id == "torch"
+                            and attr.value.attr == "_library"
+                            and attr.attr in triton_func_names
+                        ):
+                            if node.args and isinstance(node.args[0], ast.Name):
+                                self.triton_kernels.append(node.args[0].id)
+                        elif (
+                            isinstance(attr.value.value, ast.Attribute)
+                            and isinstance(attr.value.value.value, ast.Name)
+                            and attr.value.value.value.id == "torch"
+                            and attr.value.value.attr == "ops"
+                        ):
+                            self.called_functions.append(
+                                f"{attr.value.attr}::{attr.attr}"
+                            )
+                # Catch capture_triton, wrap_triton that's been
+                # imported directly
+                elif isinstance(node.func, ast.Name):
+                    if node.func.id in triton_func_names:
+                        if node.args and isinstance(node.args[0], ast.Name):
+                            self.triton_kernels.append(node.args[0].id)
+                    else:
+                        # track regular function calls for recursive analysis
+                        self.called_functions.append(node.func.id)
+
+                self.generic_visit(node)
+
+        collector = Visitor()
+        collector.visit(tree)
+
+        def extract_names_from_expr(expr: ast.expr) -> list[str]:
+            """Extract all Name references from an AST expression."""
+            names: list[str] = []
+
+            class NameExtractor(ast.NodeVisitor):
+                def visit_Name(self, node: ast.Name) -> None:
+                    names.append(node.id)
+
+                def visit_Call(self, node: ast.Call) -> None:
+                    # for function calls, visit the function and all args
+                    self.generic_visit(node)
+
+            NameExtractor().visit(expr)
+            return names
+
+        def resolve_to_kernel(obj: object) -> object | None:
+            """Check if obj is a triton kernel or wrapper and return the kernel."""
+            if isinstance(obj, (JITFunction, Autotuner)):
+                return obj
+            # handle wrappers that have a .fn attribute pointing to JITFunction
+            if callable(obj) and hasattr(obj, "fn"):
+                inner = obj.fn
+                if isinstance(inner, JITFunction):
+                    return inner
+            return None
+
+        def build_namespace(func_obj: object) -> dict[str, Any]:
+            """Build a combined namespace from a function's globals and closures."""
+            # unwrap decorated fns (e.g., @lru_cache)
+            if callable(func_obj):
+                try:
+                    func_obj = inspect.unwrap(func_obj)
+                except ValueError:
+                    pass
+            if not callable(func_obj) or not hasattr(func_obj, "__code__"):
+                return {}
+            func_closure_vars = inspect.getclosurevars(func_obj)
+            namespace: dict[str, Any] = {}
+            namespace.update(func_closure_vars.builtins)
+            namespace.update(func_closure_vars.globals)
+            namespace.update(func_closure_vars.nonlocals)
+            if hasattr(func_obj, "__globals__"):
+                namespace.update(func_obj.__globals__)
+            return namespace
+
+        all_names = build_namespace(fn)
+
+        def resolve_names_to_kernels(
+            names: list[str],
+            namespace: dict[str, Any],
+            assignments: dict[str, list[ast.expr]] | None = None,
+            visited: set[str] | None = None,
+        ) -> list[object]:
+            """
+            Resolve a list of names to triton kernels using the given namespace.
+            """
+            if visited is None:
+                visited = set()
+
+            results: list[object] = []
+            for name in names:
+                if name in visited:
+                    continue
+                visited.add(name)
+
+                if name in namespace:
+                    obj = namespace[name]
+                    kernel = resolve_to_kernel(obj)
+                    if kernel is not None:
+                        results.append(kernel)
+                        continue
+                    # recurse into callable objects (factory fn's),
+                    # unwrapping decorators if applicable
+                    if callable(obj):
+                        try:
+                            unwrapped = inspect.unwrap(obj)
+                        except ValueError:
+                            unwrapped = obj
+                        if hasattr(unwrapped, "__code__"):
+                            nested = find_triton_kernels(
+                                unwrapped, visited_fns, depth + 1
+                            )
+                            if nested:
+                                results.extend(nested)
+                                continue
+                    logger.debug("failed to resolve %s to a triton kernel", name)
+                elif assignments is not None and name in assignments:
+                    # trace through local assignments
+                    for rhs_expr in assignments[name]:
+                        referenced = extract_names_from_expr(rhs_expr)
+                        traced = resolve_names_to_kernels(
+                            referenced, namespace, assignments, visited
+                        )
+                        results.extend(traced)
+                else:
+                    logger.debug("%s not found in namespace or assignments", name)
+
+            return results
+
+        # resolve kernel names, tracing through local variables if needed
+        resolved: list[object] = []
+        seen_ids: set[int] = set()
+
+        names_to_resolve: list[str] = list(collector.triton_kernels)
+        for expr in collector.return_exprs:
+            names_to_resolve.extend(extract_names_from_expr(expr))
+
+        for name in names_to_resolve:
+            traced_objects = resolve_names_to_kernels(
+                [name], all_names, collector.assignments
+            )
+            for obj in traced_objects:
+                obj_id = id(obj)
+                if obj_id not in seen_ids:
+                    seen_ids.add(obj_id)
+                    resolved.append(obj)
+
+        for func_name in collector.called_functions:
+            func_obj = all_names.get(func_name)
+
+            if func_obj is None:
+                from torch._library.custom_ops import OPDEFS
+
+                if func_name in OPDEFS:
+                    func_obj = OPDEFS[func_name]._abstract_fn
+
+            # skip if not a callable or if it's a triton kernel itself
+            if func_obj is None or not callable(func_obj):
+                continue
+
+            # skip built-in functions and C extensions (they can't contain triton kernels)
+            if not hasattr(func_obj, "__code__"):
+                continue
+
+            try:
+                nested_kernels = find_triton_kernels(func_obj, visited_fns, depth + 1)
+                for kernel in nested_kernels:
+                    kernel_id = id(kernel)
+                    if kernel_id not in seen_ids:
+                        seen_ids.add(kernel_id)
+                        resolved.append(kernel)
+            except Exception:
+                logger.debug(
+                    "failed to analyze called function %s", func_name, exc_info=True
+                )
+
+        return resolved
+
+    return find_triton_kernels(fn)
+
+
+@exposed_in("torch.library")
+def triton_op(
+    name: str,
+    fn: Callable | None = None,
+    /,
+    *,
+    mutates_args: str | Iterable[str],
+    schema: str | None = None,
+) -> Callable:
+    """Create a custom operator whose implementation is backed by 1+ triton kernels.
+
+    This is a more structured way of using triton kernels with PyTorch.
+    Prefer using triton kernels with no ``torch.library`` custom operator wrappers
+    (like :func:`torch.library.custom_op`, :func:`torch.library.triton_op`) because
+    that is simpler;
+    only use :func:`torch.library.custom_op`/:func:`torch.library.triton_op` if you
+    want to create an operator that behaves like PyTorch built-in operators.
+    For example, you may use a ``torch.library`` wrapper API to define the
+    behavior of the triton kernel when passed a tensor subclass or under
+    a TorchDispatchMode.
+
+    Use :func:`torch.library.triton_op` instead of :func:`torch.library.custom_op`
+    when the implementation
+    consists of 1+ triton kernels. :func:`torch.library.custom_op` treats
+    custom operators as opaque (:func:`torch.compile` and
+    :func:`torch.export.export` will never trace into them), but ``triton_op``
+    makes the implementation visible to these subsystems, allowing them
+    to optimize the triton kernel(s).
+
+    Note that ``fn`` must only consist of calls to PyTorch-understood
+    operators and triton kernels. Any triton kernels called inside ``fn``
+    must be wrapped in a call to :func:`torch.library.wrap_triton`.
+
+    Args:
+        name (str): A name for the custom op that looks like "{namespace}::{name}",
+            e.g. "mylib::my_linear". The name is used as the op's stable identifier
+            in PyTorch subsystems (e.g. torch.export, FX graphs).
+            To avoid name collisions, please use your project name as the namespace;
+            e.g. all custom ops in pytorch/fbgemm use "fbgemm" as the namespace.
+        mutates_args (Iterable[str] or "unknown"): The names of args that the function mutates.
+            This MUST be accurate, otherwise, the behavior is undefined. If "unknown",
+            it pessimistically assumes that all inputs to the operator are being mutated.
+        schema (str | None): A schema string for the operator. If None
+            (recommended) we'll infer a schema for the operator from its type
+            annotations. We recommend letting us infer a schema unless you
+            have a specific reason not to.
+            Example: "(Tensor x, int y) -> (Tensor, Tensor)".
+
+    Example::
+
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> import torch
+        >>> from torch.library import triton_op, wrap_triton
+        >>>
+        >>> import triton
+        >>> from triton import language as tl
+        >>>
+        >>> @triton.jit
+        >>> def add_kernel(
+        >>>     in_ptr0,
+        >>>     in_ptr1,
+        >>>     out_ptr,
+        >>>     n_elements,
+        >>>     BLOCK_SIZE: "tl.constexpr",
+        >>> ):
+        >>>     pid = tl.program_id(axis=0)
+        >>>     block_start = pid * BLOCK_SIZE
+        >>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        >>>     mask = offsets < n_elements
+        >>>     x = tl.load(in_ptr0 + offsets, mask=mask)
+        >>>     y = tl.load(in_ptr1 + offsets, mask=mask)
+        >>>     output = x + y
+        >>>     tl.store(out_ptr + offsets, output, mask=mask)
+        >>>
+        >>> @triton_op("mylib::add", mutates_args={})
+        >>> def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        >>>     output = torch.empty_like(x)
+        >>>     n_elements = output.numel()
+        >>>
+        >>>     def grid(meta):
+        >>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+        >>>
+        >>>     # NB: we need to wrap the triton kernel in a call to wrap_triton
+        >>>     wrap_triton(add_kernel)[grid](x, y, output, n_elements, 16)
+        >>>     return output
+        >>>
+        >>> @torch.compile
+        >>> def f(x, y):
+        >>>     return add(x, y)
+        >>>
+        >>> x = torch.randn(3, device="cuda")
+        >>> y = torch.randn(3, device="cuda")
+        >>>
+        >>> z = f(x, y)
+        >>> assert torch.allclose(z, x + y)
+
+    """
+
+    def dec(fn: Callable[..., object]) -> CustomOpDef:
+        def backend_fn(*args, **kwargs):  # type: ignore[no-untyped-def]
+            # Optimization: we're passing regular Tensors into the triton kernel, so
+            # no need to go through HOP dispatch
+            with set_wrap_triton_enabled(False):
+                return fn(*args, **kwargs)
+
+        result = custom_op(
+            name,
+            backend_fn,
+            mutates_args=mutates_args,
+            schema=infer_schema(fn, mutates_args=mutates_args),
+        )
+        from .._subclasses.functional_tensor import FunctionalTensorMode
+
+        # We require that the user pass us a function that is make_fx traceable,
+        # so we can just register it as the Fake/meta kernel.
+        result.register_fake(fn)
+
+        # We decompose the operator when FunctionalTensorMode is active.
+        # The goal is to decompose the operator in AOTDispatcher.
+        # - With torch.compile, this means that the backend (usually Inductor)
+        #   can see a call to the triton kernel(s) and so it can directly optimize
+        #   them by inlining them into the lowering process.
+        def functional_decomp(  # type: ignore[no-untyped-def]
+            mode, op, types, args, kwargs
+        ):
+            # NOTE [Export custom triton op]
+            # For torch.export (strict and non-strict), we don't do functional decomposition.
+            # Instead, we preserve the custom triton ops as custom ops. This is because we want
+            # the exported program to be high-level and serializable. If we decompose
+            # the custom op to a functional hop and make it a node in exported program,
+            # we need to figure out ways of serializing the hop and its arguments, which can be triton.jited
+            # functions and triton dtypes. This is undesirable because:
+            # - it can be tedious to maintain a layer that serializes the jited function (e.g. with a string) and dtypes.
+            # - exported program will contain the implementation detail (e.g. triton source code) for a specific
+            #   backend (GPU), which is probably at a wrong level of abstraction.
+            # - changes to triton or the serialization logic for triton arguments can be BC breaking
+            #
+            # In the short term, we expect users to have a separate aot_compile stage that compiles the exported program
+            # into a Cubin file on the same machine that users call export, which does autotuning and removes triton
+            # dependency and serve the model with Cubin. This guarantees that triton changes won't break BC.
+            # In the long term, we may export multiple cubins for the triton op directly
+            from torch.export._trace import custom_triton_ops_decomposition_disabled
+
+            if custom_triton_ops_decomposition_disabled():
+                return mode.__torch_dispatch__(op, types, args, kwargs)
+            else:
+                # TODO: https://github.com/pytorch/pytorch/issues/160333
+                # We should deduplicate the unrecognized_types logic.
+                import torch._subclasses
+
+                unrecognized_types = [
+                    t
+                    for t in types
+                    if not issubclass(t, torch._subclasses.FakeTensor)
+                    and t
+                    not in [
+                        torch.Tensor,
+                        torch._subclasses.functional_tensor.FunctionalTensor,
+                    ]
+                ]
+
+                if unrecognized_types:
+                    return NotImplemented
+                with mode:
+                    return fn(*args, **kwargs)
+
+        triton_kernels = get_inner_triton_kernels(fn)
+        triton_ops_to_kernels[name] = triton_kernels
+        result.register_torch_dispatch(FunctionalTensorMode, functional_decomp)
+        return result
+
+    if fn is None:
+        return dec
+    else:
+        return dec(fn)
+
+
+wrap_triton_enabled = threading.local()
+wrap_triton_enabled_default = True
+
+
+@contextlib.contextmanager
+def set_wrap_triton_enabled(enabled: bool) -> Generator[None, None, None]:
+    """If triton kernels annotated with @wrap_triton should dispatch via HOP
+    or go straight to the triton kernel execution.
+
+    We have this switch because eager-mode performance of HOP dispatch is slow
+    enough to matter (~1ms) and we know that wrap_triton isn't necessary in
+    some situations (eager-mode with regular Tensors)
+    """
+    try:
+        prev = is_wrap_triton_enabled()
+        wrap_triton_enabled.value = enabled
+        yield
+    finally:
+        wrap_triton_enabled.value = prev
+
+
+def is_wrap_triton_enabled() -> bool:
+    return getattr(wrap_triton_enabled, "value", wrap_triton_enabled_default)
+
+
+def capture_triton(triton_kernel: Callable, /) -> Any:
+    """This API has been renamed to wrap_triton"""
+    return wrap_triton(triton_kernel)
+
+
+@exposed_in("torch.library")
+def wrap_triton(triton_kernel: Callable, /) -> Any:
+    """Allows capture of a triton kernel into a graph via make_fx or
+    non-strict ``torch.export``.
+
+    These technologies perform Dispatcher-based tracing (via
+    ``__torch_dispatch__``) and cannot see calls to raw triton kernels.
+    The ``wrap_triton`` API wraps a triton kernel into a callable that
+    can actually be traced into a graph.
+
+    Please use this API together with :func:`torch.library.triton_op`.
+
+    Examples:
+
+        >>> # xdoctest: +SKIP
+        >>> import torch
+        >>> import triton
+        >>> from triton import language as tl
+        >>> from torch.fx.experimental.proxy_tensor import make_fx
+        >>> from torch.library import wrap_triton
+        >>>
+        >>> @triton.jit
+        >>> def add_kernel(
+        >>>     in_ptr0,
+        >>>     in_ptr1,
+        >>>     out_ptr,
+        >>>     n_elements,
+        >>>     BLOCK_SIZE: "tl.constexpr",
+        >>> ):
+        >>>     pid = tl.program_id(axis=0)
+        >>>     block_start = pid * BLOCK_SIZE
+        >>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        >>>     mask = offsets < n_elements
+        >>>     x = tl.load(in_ptr0 + offsets, mask=mask)
+        >>>     y = tl.load(in_ptr1 + offsets, mask=mask)
+        >>>     output = x + y
+        >>>     tl.store(out_ptr + offsets, output, mask=mask)
+        >>>
+        >>> def add(x, y):
+        >>>     output = torch.empty_like(x)
+        >>>     n_elements = output.numel()
+        >>>
+        >>>     def grid_fn(meta):
+        >>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+        >>>
+        >>>     wrap_triton(add_kernel)[grid_fn](x, y, output, n_elements, 16)
+        >>>     return output
+        >>>
+        >>> x = torch.randn(3, device="cuda")
+        >>> y = torch.randn(3, device="cuda")
+        >>> gm = make_fx(add)(x, y)
+        >>> print(gm.code)
+        >>> # def forward(self, x_1, y_1):
+        >>> #     empty_like = torch.ops.aten.empty_like.default(x_1, pin_memory = False)
+        >>> #     triton_kernel_wrapper_mutation_proxy = triton_kernel_wrapper_mutation(
+        >>> #         kernel_idx = 0, constant_args_idx = 0,
+        >>> #         grid = [(1, 1, 1)], kwargs = {
+        >>> #             'in_ptr0': x_1, 'in_ptr1': y_1, 'out_ptr': empty_like,
+        >>> #             'n_elements': 3, 'BLOCK_SIZE': 16
+        >>> #         })
+        >>> #     return empty_like
+
+    """
+    from triton.runtime.autotuner import Autotuner
+    from triton.runtime.jit import JITFunction
+
+    from torch._higher_order_ops.triton_kernel_wrap import TraceableTritonKernelWrapper
+
+    if not isinstance(triton_kernel, (JITFunction, Autotuner)):
+        raise RuntimeError(
+            "wrap_triton only works on functions annotated with triton.jit or triton.autotune"
+        )
+    if not is_wrap_triton_enabled():
+        return triton_kernel
+    return TraceableTritonKernelWrapper(triton_kernel, None, None)

kronos-venv/lib/python3.12/site-packages/torch/_library/utils.py

@@ -0,0 +1,662 @@
+# mypy: allow-untyped-defs
+import dataclasses
+import inspect
+import sys
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any, Literal, overload
+
+import torch
+import torch.utils._pytree as pytree
+import torchgen
+from torch import _C, _utils_internal
+from torch._ops import OpOverload
+
+
+@dataclasses.dataclass
+class Kernel:
+    """Models a (function, source location)"""
+
+    func: Callable
+    source: str
+
+    def __call__(self, *args, **kwargs):
+        return self.func(*args, **kwargs)
+
+
+class RegistrationHandle:
+    """Does something when someone calls .destroy() on it"""
+
+    def __init__(self, on_destroy: Callable):
+        self._on_destroy = on_destroy
+
+    def destroy(self) -> None:
+        self._on_destroy()
+
+
+def get_source(stacklevel: int) -> str:
+    """Get a string that represents the caller.
+
+    Example: "/path/to/foo.py:42"
+
+    Use stacklevel=1 to get the caller's source
+    Use stacklevel=2 to get the caller's caller's source
+    etc.
+    """
+    frame = inspect.getframeinfo(sys._getframe(stacklevel))
+    source = f"{frame.filename}:{frame.lineno}"
+    return source
+
+
+def parse_namespace(qualname: str) -> tuple[str, str]:
+    splits = qualname.split("::")
+    if len(splits) != 2:
+        raise ValueError(
+            f"Expected `qualname` to be of the form "
+            f'"namespace::name", but got {qualname}. '
+            f"The qualname passed to the torch.library APIs must consist "
+            f"of a namespace and a name, e.g. aten::sin"
+        )
+    return splits[0], splits[1]
+
+
+def lookup_op(qualname: str) -> OpOverload:
+    namespace, name = parse_namespace(qualname)
+    if "." in name:
+        name, overload = name.split(".")
+    else:
+        overload = "default"
+    ns = getattr(torch.ops, namespace)
+    packet = getattr(ns, name)
+    return getattr(packet, overload)
+
+
+def is_builtin(op: OpOverload) -> bool:
+    if not isinstance(op, OpOverload):
+        raise AssertionError(f"op must be OpOverload, got {type(op)}")
+    return op.namespace in {"aten", "prim", "prims"}
+
+
+def is_functional_schema(schema: Any, *, allow_valid_view: bool = False) -> bool:
+    """Check if the schema is functional.
+
+    An operator is functional if:
+    - it does not mutate any of its inputs
+    - If no view are allowed
+        - it does not return a view on any of its inputs
+    - If valid views are allowed
+        - it is not a view or a view with a single input Tensor and single output Tensor
+    - it has at least one return
+    """
+
+    def is_functional(schema):
+        if schema.is_mutable:
+            return False
+        rets = schema.returns
+        is_non_mutating_view = len(rets) > 0 and any(
+            r.alias_info is not None and not r.alias_info.is_write for r in rets
+        )
+        num_tensor_inputs = 0
+        num_tensor_outputs = 0
+
+        if isinstance(schema, torch.FunctionSchema):
+            for arg in schema.arguments:
+                if isinstance(arg.type, torch.TensorType):
+                    num_tensor_inputs += 1
+
+            for ret in schema.returns:
+                if isinstance(ret.type, torch.TensorType):
+                    num_tensor_outputs += 1
+
+        elif isinstance(schema, torchgen.model.FunctionSchema):
+            for argument in schema.arguments.flat_non_out:
+                if argument.type.is_tensor_like():
+                    num_tensor_inputs += 1
+
+            for ret_arg in schema.returns:
+                if ret_arg.type.is_tensor_like():
+                    num_tensor_outputs += 1
+
+        if is_non_mutating_view:
+            return allow_valid_view and (
+                num_tensor_inputs == 1 and num_tensor_outputs == 1
+            )
+        if not schema.returns:
+            return False
+        return True
+
+    if isinstance(schema, torch._C.FunctionSchema):
+        return is_functional(schema)
+
+    # Lazy import because not all PyTorch builds have torchgen
+    from torchgen.model import FunctionSchema
+
+    if isinstance(schema, str):
+        schema = FunctionSchema.parse(schema)
+    if not isinstance(schema, FunctionSchema):
+        raise AssertionError(f"schema must be FunctionSchema, got {type(schema)}")
+    return is_functional(schema)
+
+
+# should be torch._C.JitType but that annotation is busted
+def is_tensorlist_like_type(typ: Any) -> bool:
+    return (
+        typ == _C.ListType(_C.TensorType.get())
+        or typ == _C.ListType(_C.OptionalType(_C.TensorType.get()))
+        or typ == _C.OptionalType(_C.ListType(_C.TensorType.get()))
+        or typ == _C.OptionalType(_C.ListType(_C.OptionalType(_C.TensorType.get())))
+    )
+
+
+# should be torch._C.JitType but that annotation is busted
+def is_tensor_like_type(typ: Any) -> bool:
+    return typ == _C.TensorType.get() or typ == _C.OptionalType(_C.TensorType.get())
+
+
+def mutates_and_returns_first_arg(op: OpOverload):
+    """Check if an op is an inplace aten op, i.e. it mutates and returns the first arg.
+
+    TODO: torchgen/model.py's FunctionSchema.parse is the source of truth for this,
+    but not all PyTorch builds have torchgen (due to the yaml dependency being weird).
+    Figure this out.
+
+    Example: add_(Tensor(a!) x, Tensor y) -> Tensor(a)
+    """
+    if op.namespace != "aten":
+        return False
+    schema = op._schema
+    if len(schema.returns) != 1:
+        return False
+    if schema.returns[0].alias_info is None:
+        return False
+    alias_set = schema.returns[0].alias_info.after_set
+    if len(alias_set) != 1:
+        return False
+    loc = next(iter(alias_set))
+    if len(schema.arguments) < 1:
+        return False
+    first_arg = schema.arguments[0]
+    if first_arg.alias_info is None:
+        return False
+    if not first_arg.alias_info.is_write:
+        return False
+    alias_set = first_arg.alias_info.after_set
+    if len(alias_set) != 1:
+        return False
+    if loc != next(iter(alias_set)):
+        return False
+    for arg in schema.arguments[1:]:
+        if arg.alias_info is not None:
+            return False
+    return True
+
+
+def fill_defaults(schema, args, kwargs):
+    new_args = []
+    new_kwargs = {}
+    for i in range(len(schema.arguments)):
+        info = schema.arguments[i]
+        if info.kwarg_only:
+            if info.name in kwargs:
+                new_kwargs[info.name] = kwargs[info.name]
+            else:
+                new_kwargs[info.name] = info.default_value
+        else:
+            if i < len(args):
+                new_args.append(args[i])
+            else:
+                new_args.append(info.default_value)
+    return tuple(new_args), new_kwargs
+
+
+def zip_schema(
+    schema: _C.FunctionSchema, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> Iterable[tuple[_C.Argument, Any]]:
+    """zips schema.arguments and (args, kwargs) together.
+
+    Assumes that (args, kwargs) were the inputs to some torch._ops.OpOverload:
+    that is, (args, kwargs) must be bindable to the schema (args, kwargs).
+    """
+    if len(schema.arguments) < len(args) + len(kwargs):
+        raise AssertionError(
+            f"schema has {len(schema.arguments)} arguments but got {len(args)} args and {len(kwargs)} kwargs"
+        )
+    for i in range(len(schema.arguments)):
+        info = schema.arguments[i]
+        if info.kwarg_only:
+            if info.name in kwargs:
+                yield info, kwargs[info.name]
+            continue
+        if i >= len(args):
+            if not info.kwarg_only and info.name in kwargs:
+                yield info, kwargs[info.name]
+            # args that are equal to their default values are not populated
+            # if they are followed by args that are equal to their defaults.
+            # Skip these.
+            continue
+        yield info, args[i]
+    return
+
+
+def hop_schema_from_fx_node(node):
+    from torchgen.gen_schema_utils import FunctionSchemaGen
+
+    hop = node.target
+    if not isinstance(hop, torch._ops.HigherOrderOperator):
+        raise RuntimeError("fx_node's target must be a hop.")
+
+    def _collect_example_val(node):
+        meta_val = node.meta.get("val", None)
+        if meta_val is None:
+            if node.op != "get_attr":
+                raise AssertionError(
+                    f"node.op must be 'get_attr' when val is None, got {node.op!r}"
+                )
+            meta_val = getattr(node.graph.owning_module, node.target)
+        return meta_val
+
+    example_inputs = []
+    for arg in node.args:
+        if isinstance(arg, (torch.fx.Node, torch.fx.node.Node)):
+            example_inputs.append(_collect_example_val(arg))
+        elif isinstance(
+            arg, (torch.fx.immutable_collections.immutable_list, list, tuple)
+        ):
+            example_inputs.append([_collect_example_val(x) for x in arg])
+        else:
+            raise RuntimeError(f"Unsupported arg type {type(arg)}")
+
+    # Bound the arguments to make sure number of inputs are correct
+    bound_args: inspect.BoundArguments = inspect.signature(hop.__call__).bind(
+        *example_inputs
+    )
+
+    # We treat example_output as a single value in return. This is to differentiate 1. return a single val
+    # vs 2. return a tuple with one element.
+    example_output = _collect_example_val(node)
+    return FunctionSchemaGen.from_example(
+        hop._name, tuple(bound_args.arguments.items()), (list(example_output),)
+    )
+
+
+def can_generate_trivial_fake_impl(op: OpOverload) -> bool:
+    if not isinstance(op, OpOverload):
+        raise AssertionError(f"op must be OpOverload, got {type(op)}")
+    if is_builtin(op):
+        # We control the built-ins. These may (in rare cases)
+        # do input metadata mutation (which we have banned on custom ops)
+        return False
+    schema = op._schema
+    # It's suspicious if the op is not mutable but returns nothing, so we return False out of an abundance of caution
+    if not schema.is_mutable:
+        return False
+    if len(schema.returns) > 0:
+        return False
+    # If the op returns nothing, then it has a trivial fake impl.
+    return True
+
+
+def requires_set_python_module() -> bool:
+    """If an op was defined in C++ and extended from Python using the
+    torch.library APIs, returns if we require that there have been a
+    m.set_python_module("mylib.ops") call from C++ that associates
+    the C++ op with a python module.
+    """
+    return getattr(_utils_internal, "REQUIRES_SET_PYTHON_MODULE", True)
+
+
+def handle_dispatch_mode(curr_mode, op_overload, *args, **kwargs):
+    if not isinstance(curr_mode, torch.utils._python_dispatch.TorchDispatchMode):
+        raise AssertionError(
+            f"curr_mode must be TorchDispatchMode, got {type(curr_mode)}"
+        )
+    args_flattened, _ = torch.utils._pytree.tree_flatten((args, kwargs.values()))
+    # TODO: need to double check the semantics of the "types" argument to torch_dispatch.
+    # It's generated in PyInterpreter.cpp, but seems to be generated in two places,
+    # where in one case we only include tensors with the python key, and in another
+    # we include **all** tensors.
+    overload_types = [
+        type(a)
+        for a in args_flattened
+        if isinstance(a, torch.Tensor)
+        and torch._C._dispatch_keys(a).has(torch._C.DispatchKey.Python)
+    ]
+    # TODO: check that I got these args correct (in C++, we pass in "0000"??)
+
+    return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs)
+
+
+def has_kwarg_only_args(schema: _C.FunctionSchema):
+    return any(a.kwarg_only for a in schema.arguments)
+
+
+def has_kwarg_only_tensors(schema: _C.FunctionSchema):
+    for a in schema.arguments:
+        if not (is_tensor_like_type(a.type) or is_tensorlist_like_type(a.type)):
+            continue
+        if not a.kwarg_only:
+            continue
+        return True
+    return False
+
+
+def has_tensor_arg(schema: _C.FunctionSchema) -> bool:
+    """
+    Given a schema, returns True if the schema has a Tensor arg.
+    A Tensor arg is any arg with a type annotation that might involve Tensor.
+    """
+    return any(
+        (is_tensor_like_type(a.type) or is_tensorlist_like_type(a.type))
+        for a in schema.arguments
+    )
+
+
+def get_device_arg_index(schema: _C.FunctionSchema) -> int | None:
+    """
+    Given a schema, returns the id of the `device: torch.device` argument.
+    If it does not exist, returns None.
+    """
+    for index, arg in enumerate(schema.arguments):
+        if arg.type is _C.DeviceObjType.get() and arg.name == "device":
+            return index
+    return None
+
+
+def iter_tensors(
+    args: tuple[Any], kwargs: dict[str, Any], allowed_nesting: int = 1
+) -> Iterator[torch.Tensor]:
+    def check(arg):
+        if isinstance(arg, torch.Tensor):
+            yield arg
+        elif allowed_nesting > 0 and isinstance(arg, (tuple, list)):
+            yield from iter_tensors(tuple(arg), {}, allowed_nesting - 1)
+
+    for arg in args:
+        yield from check(arg)
+    for kwarg in kwargs.values():
+        yield from check(kwarg)
+
+
+def check_aliasing_constraint(name, prev, result, get_module=lambda: "???"):
+    """
+    custom operators' outputs must not alias any inputs or other outputs.
+    """
+    storages = {t.untyped_storage()._cdata for t in prev if isinstance(t, torch.Tensor)}
+    tuple_result = result
+    if not isinstance(result, tuple):
+        tuple_result = (result,)
+    for tensor in iter_tensors(tuple_result, {}):
+        key = tensor.untyped_storage()._cdata
+        if tensor.untyped_storage()._cdata in storages:
+            raise RuntimeError(
+                f"{name} (with implementation in {get_module()}): "
+                f"The output of this custom operator (1) must not "
+                f"also be an input to this custom operator and "
+                f"(2) may not alias any inputs to this custom operator "
+                f"or other returns. "
+                f"The most common way to trigger this error is if "
+                f"we have y = custom_op(x) and y and x are the same Tensor. "
+                f"Please instead return a clone of the offending output "
+                f"tensor(s) (e.g. return x.clone()) or refactor the custom "
+                f"operator to not return y."
+            )
+        storages.add(key)
+
+
+def _c_check_aliasing_constraint(name, args, kwargs, result, get_module=lambda: "???"):
+    """
+    custom operators' outputs must not have any aliases
+    This version uses C++ implementation for perf.
+    Only List container is supported.
+    Tensors in Lists with not only Tensors are checked.
+    """
+    tuple_result = result
+    if not isinstance(result, tuple):
+        tuple_result = (result,)
+    if _C._any_output_is_alias_to_input_or_output(args, kwargs, tuple_result):
+        raise RuntimeError(
+            f"{name} (with implementation in {get_module()}): "
+            f"The output of this custom operator (1) must not "
+            f"also be an input to this custom operator and "
+            f"(2) may not alias any inputs to this custom operator "
+            f"or other returns. "
+            f"The most common way to trigger this error is if "
+            f"we have y = custom_op(x) and y and x are the same Tensor. "
+            f"Please instead return a clone of the offending output "
+            f"tensor(s) (e.g. return x.clone()) or refactor the custom "
+            f"operator to not return y."
+        )
+
+
+class MutationChecker:
+    """
+    Check if an operator mutated its arguments.
+    Usage:
+
+    checker = MutationChecker(op, flat_args, args_spec)
+    op(*args, **kwargs)
+    checker.check()
+    """
+
+    def __init__(self, op, flat_args, args_spec):
+        self.op = op
+        self.args_spec = args_spec
+        self.flat_args = flat_args
+        self.real_pre_hashes = [
+            hash_tensor(a) if isinstance(a, torch.Tensor) else None for a in flat_args
+        ]
+
+    def check(self):
+        real_post_hashes = [
+            hash_tensor(a) if isinstance(a, torch.Tensor) else None
+            for a in self.flat_args
+        ]
+        was_mutated = [
+            not torch.equal(pre, post)
+            and not (pre.isnan().all() and post.isnan().all())
+            if isinstance(pre, torch.Tensor) and isinstance(post, torch.Tensor)
+            else None
+            for pre, post in zip(self.real_pre_hashes, real_post_hashes)
+        ]
+        was_mutated_args, was_mutated_kwargs = pytree.tree_unflatten(
+            was_mutated, self.args_spec
+        )
+        for info, was_mutated in zip_schema(
+            self.op._schema, was_mutated_args, was_mutated_kwargs
+        ):
+
+            def check_one(info, was_mutated):
+                if info.is_write == was_mutated:
+                    return
+                raise RuntimeError(
+                    f"{self.op._name}: for argument '{info.name}': the operator's schema "
+                    f"{self.op._schema} specified that "
+                    f"the operator {'mutates' if info.is_write else 'does not mutate'} "
+                    f"the argument, but this seems to be empirically wrong. "
+                    f"Please make the schema and operator behavior consistent. "
+                    f"You can specify that an operator mutates a Tensor by "
+                    f"e.g. changing its schema type from 'Tensor name' to 'Tensor(a!) name'"
+                    f"(use different identifiers (a, b, c, ...) for different Tensors)"
+                )
+
+            if is_tensor_like_type(info.type):
+                check_one(info, was_mutated)
+            elif is_tensorlist_like_type(info.type):
+                was_any_mutated = False if was_mutated is None else any(was_mutated)
+                check_one(info, was_any_mutated)
+
+
+def hash_tensor(t: torch.Tensor) -> torch.Tensor:
+    """Some inexpensive hash. Used as a quick and dirty indicator for tensor mutation"""
+    return t.detach().float().mean()
+
+
+def has_fake_kernel(op: torch._ops.OpOverload) -> bool:
+    """If an operator (that stays alive until FakeTensorMode) has a Fake kernel.
+    Don't use this if the operator decomposes before FakeTensorMode.
+    """
+    if can_generate_trivial_fake_impl(op):
+        return True
+    name = op._name
+    if torch._C._dispatch_has_kernel_for_dispatch_key(
+        name, "CompositeImplicitAutograd"
+    ):
+        return True
+    opdef = torch._library.custom_ops._maybe_get_opdef(name)
+    if opdef is None:
+        # the non-torch.library.custom_op path
+        if torch._C._dispatch_has_kernel_for_dispatch_key(
+            name, "CompositeExplicitAutograd"
+        ):
+            return True
+        entry = torch._library.simple_registry.singleton.find(name)
+        if entry.fake_impl.kernel is not None:
+            return True
+        if torch._C._dispatch_has_kernel_for_dispatch_key(name, "Meta"):
+            return True
+    else:
+        # the torch.library.custom_op path
+        if opdef._abstract_fn is not None:
+            return True
+    return False
+
+
+def mutated_args_kwargs(schema: _C.FunctionSchema) -> tuple[list[int], list[str]]:
+    idxs = []
+    keys = []
+    for i, info in enumerate(schema.arguments):
+        if info.alias_info is not None and info.alias_info.is_write:
+            if info.kwarg_only:
+                keys.append(info.name)
+            else:
+                idxs.append(i)
+    return idxs, keys
+
+
+tags_by_priority = [
+    _C.Tag.needs_exact_strides,
+    _C.Tag.needs_contiguous_strides,
+    _C.Tag.needs_fixed_stride_order,
+    _C.Tag.flexible_layout,
+]
+
+
+# Case 1: with_default=True (or omitted). Return type is guaranteed to be a Tag.
+@overload
+def get_layout_constraint_tag(
+    fn: Any, *, with_default: Literal[True] = True
+) -> _C.Tag: ...
+
+
+# Case 2: with_default=False. Return type can be a Tag or None.
+@overload
+def get_layout_constraint_tag(
+    fn: Any, *, with_default: Literal[False]
+) -> _C.Tag | None: ...
+
+
+def get_layout_constraint_tag(fn, *, with_default=True):
+    for tag in tags_by_priority:
+        if tag in fn.tags:
+            return tag
+    if with_default:
+        if is_builtin(fn):
+            return _C.Tag.flexible_layout
+        import torch._functorch
+        from torch._functorch import config
+
+        return getattr(torch._C.Tag, config.custom_op_default_layout_constraint)
+    return None
+
+
+# List of random functions that should be treated as impure
+_RANDOM_FUNCTIONS = {
+    torch.rand,
+    torch.randn,
+    torch.randint,
+    torch.randperm,
+    torch.rand_like,
+    torch.randn_like,
+    torch.randint_like,
+    torch.normal,
+    torch.poisson,
+    torch.bernoulli,
+    torch.multinomial,
+}
+
+
+def is_impure(
+    op: Callable,
+    *,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    impure_random: bool = True,
+) -> bool:
+    """
+    An operator is impure if it:
+    - Mutates its inputs (has a mutable schema)
+    - Has nondeterministic/random behavior that mutates RNG state
+    - Is explicitly marked as effectful via torch.library._register_effectful_op
+
+    Args:
+        op: The operator to check (function, OpOverload, HigherOrderOperator, etc.)
+        args: Optional arguments that would be passed to the callable
+        kwargs: Optional keyword arguments that would be passed to the callable
+        impure_random: Whether to treat random operations as impure (default: True)
+
+    Returns:
+        bool: True if the callable has side effects, False otherwise
+    """
+    # Import here to avoid circular dependencies
+    from torch._higher_order_ops.effects import _get_effect
+    from torch.fx.node import _side_effectful_functions
+
+    if isinstance(op, torch._ops.OpOverload):
+        schema = getattr(op, "_schema", None)
+        if schema is not None and schema.is_mutable:
+            return True
+
+        if op in _side_effectful_functions:
+            return True
+
+        if _get_effect(op) is not None:
+            return True
+
+    if isinstance(op, torch._ops.HigherOrderOperator):
+        if op in (
+            torch.ops.higher_order.auto_functionalized,
+            torch.ops.higher_order.auto_functionalized_v2,
+        ):
+            # Check if the auto-functionalized operator (the first argument) is
+            # side-effectful
+            if args and len(args) > 0:
+                return args[0] in _side_effectful_functions
+
+        if _get_effect(op) is not None:
+            return True
+
+        if op in _side_effectful_functions:
+            return True
+
+        return False
+
+    # Impure since it mutates RNG state
+    if impure_random and getattr(op, "_nondeterministic_seeded", False):
+        return True
+
+    # Handle Python random functions that don't have _nondeterministic_seeded
+    # but still affect global RNG state (issue #151524)
+    # These should be impure regardless of impure_random setting to maintain
+    # consistency between eager and compiled execution
+    # All random operations are impure to ensure consistent behavior
+    # between eager and compiled execution, regardless of generator usage
+    if op in _RANDOM_FUNCTIONS:
+        return True
+
+    schema = getattr(op, "_schema", None)
+    if schema is not None and schema.is_mutable:
+        return True
+
+    if op in _side_effectful_functions:
+        return True
+
+    return False