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.
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)