"""
Repro: tensor subclass metadata guard crashes with non-scalar tensors.

The metadata guard (guards.py:2460) does:
    x.__tensor_flatten__()[1] == original_metadata

For dict comparison, Python does element-wise value comparison. When a value
is a multi-element tensor, `tensor == tensor` returns a multi-element bool
tensor. dict.__eq__ then tries bool() on it, which raises:
    "Boolean value of Tensor with more than one value is ambiguous"

This means __tensor_flatten__ metadata CANNOT contain non-scalar tensors,
though nothing documents or enforces this.
"""

import torch

class WeightedTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data, weights):
        return torch.Tensor._make_wrapper_subclass(
            cls,
            data.shape,
            dtype=data.dtype,
            device=data.device,
            requires_grad=data.requires_grad,
        )

    def __init__(self, data, weights):
        self._data = data
        self._weights = weights

    def __tensor_flatten__(self):
        # Putting a non-scalar tensor in metadata breaks the guard
        return ["_data"], {"weights": self._weights}

    @staticmethod
    def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
        return WeightedTensor(inner_tensors["_data"], meta["weights"])

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        def unwrap(t):
            return t._data if isinstance(t, WeightedTensor) else t

        out = func(
            *torch.utils._pytree.tree_map(unwrap, args),
            **torch.utils._pytree.tree_map(unwrap, kwargs or {}),
        )
        if isinstance(out, torch.Tensor) and not isinstance(out, WeightedTensor):
            return WeightedTensor(out, args[0]._weights)
        return out


def fn(x):
    return x * 2 + 1


# Works: scalar tensor in metadata
print("=== Scalar tensor in metadata ===")
x_ok = WeightedTensor(torch.randn(4), torch.tensor(1.0))
result = torch.compile(fn, backend="eager")(x_ok)


torch._dynamo.reset()

# Crashes: multi-element tensor in metadata
print("\n=== Multi-element tensor in metadata (5 elements) ===")
x_bad = WeightedTensor(torch.randn(4), torch.randn(5))
result = torch.compile(fn, backend="eager")(x_bad)

---

AssertionError: Guard failed on the same frame it was created. This is a bug - please create an issue.Guard fail reason: 0/0: Boolean value of Tensor with more than one value is ambiguous