Support for torch.equal when exporting with Dynamo

Is there a way to export models that call torch.equal? I don’t have access to the source code to rewrite those at the moment.

When I try to export using Dynamo, I get:
Unsupported: data dependent operator: aten.equal.default; to enable, set torch._dynamo.config.capture_scalar_outputs = True

However setting torch._dynamo.config.capture_scalar_outputs = True does not seem to have any effect.

Any insight is appreciated!

What is the model with torch.equal() doing with it?

It’s using GPyTorch, which calls that function in different places. I’m mainly hitting this issue with the kernel classes where they check that the data sets are different during the covariance computation:

        if last_dim_is_batch:
            x1 = x1.transpose(-1, -2).unsqueeze(-1)
            x2 = x2.transpose(-1, -2).unsqueeze(-1)

        x1_eq_x2 = torch.equal(x1, x2)
        res = None

        if diag:
            # Special case the diagonal because we can return all zeros most of the time.
            if x1_eq_x2:
                return torch.zeros(*x1.shape[:-2], x1.shape[-2], dtype=x1.dtype, device=x1.device)
            else:
                res = torch.linalg.norm(x1 - x2, dim=-1)  # 2-norm by default
                return res.pow(2) if square_dist else res
        else:
            dist_func = sq_dist if square_dist else dist
            return dist_func(x1, x2, x1_eq_x2)

See DataDependentOutputException with aten.equal.default and Dynamo export · Issue #145970 · pytorch/pytorch · GitHub, it’s difficult to support data-dependent control flow