real tensor prop for composite ops (pytorch#135717)

Fixes pytorch#135632 Adds real tensor propagation for decompositions, checking any symbols on their outputs Pull Request resolved: pytorch#135717 Approved by: https://github.com/ezyang

Author: pianpwk

test/export/test_export.py

@@ -890,6 +890,25 @@ def forward(self, x):
  torch.allclose(ep.module()(torch.zeros(2, 3)), torch.ones(2, 3) * 21)
  )
 
+ @testing.expectedFailureTrainingIRToRunDecompNonStrict # TODO(pianpwk): user_output signature
+ def test_real_tensor_for_max_op(self):
+ class Foo(torch.nn.Module):
+ def forward(self, x, y):
+ x = x[x > 0]
+ y = y[y > 0]
+ return max(x.shape[0], y.shape[0])
+
+ model = Foo()
+ inputs = (torch.randn(64), torch.randn(64))
+ with torch._functorch.config.patch(fake_tensor_propagate_real_tensors=True):
+ ep = export(model, inputs)
+
+ self.assertEqual(ep.module()(*inputs), model(*inputs))
+ x = torch.zeros(64)
+ y = torch.ones(64)
+ self.assertEqual(ep.module()(x, x), model(x, x))
+ self.assertEqual(ep.module()(x, y), model(x, y))
+
  def test_export_script_module(self):
  class Foo(torch.nn.Module):
  def forward(self, rv: torch.Tensor, t: torch.Tensor):

torch/_subclasses/fake_tensor.py

@@ -1877,6 +1877,7 @@ def maybe_to_real_tensor(t: T) -> Optional[Union[T, Tensor]]:
  for a in flat_args
  )
  ):
+ log.debug("propagate_real_tensors %s", func)
  real_flat_args = [maybe_to_real_tensor(a) for a in flat_args]
  real_args, real_kwargs = pytree.tree_unflatten(real_flat_args, args_spec)
  real_out = func(*real_args, **real_kwargs)
@@ -1888,7 +1889,7 @@ def maybe_to_real_tensor(t: T) -> Optional[Union[T, Tensor]]:
  # However, if there's a bug in the condition above, this condition
  # will also trigger.
  log.debug(
- "propagate_real_tensors skipped %s(%s, %s) %s",
+ "SKIPPED propagate_real_tensors %s(%s, %s) %s",
  func,
  flat_arg_fake_tensors,
  flat_args,
@@ -1898,17 +1899,40 @@ def maybe_to_real_tensor(t: T) -> Optional[Union[T, Tensor]]:
  def maybe_propagate_real_tensors(fake_out: T) -> T:
  import sympy
 
+ log.debug("maybe_propagate_real_tensors %s", func)
+
  def go(t: object, real_t: Tensor) -> None:
  if isinstance(t, FakeTensor):
  # NB: unconditionally overwrite
+ log.debug(
+ "maybe_propagate_real_tensors %s -> %s", id(t), id(real_t)
+ )
  t.real_tensor = real_t
+ for s, real_s in zip(t.size(), real_t.size()):
+ go(s, real_s) # type: ignore[arg-type]
+ for s, real_s in zip(t.stride(), real_t.stride()):
+ go(s, real_s) # type: ignore[arg-type]
+ go(t.storage_offset(), real_t.storage_offset()) # type: ignore[arg-type]
  elif isinstance(t, py_sym_types) and free_unbacked_symbols(t):
  if isinstance(t.node.expr, sympy.Symbol):
  assert self.shape_env is not None
  self.shape_env.set_unbacked_var_to_val(t.node.expr, real_t)
 
  if real_out is not nil:
- tree_map_(go, fake_out, real_out)
+ if (
+ not isinstance(fake_out, Tensor)
+ and not isinstance(real_out, Tensor)
+ and type(fake_out) != type(real_out)
+ ):
+ # This can happen when decompositions have different return types,
+ # e.g. namedtuple vs. tuple vs. list.
+ tree_map_(
+ go,
+ tuple(pytree.tree_flatten(fake_out)),
+ tuple(pytree.tree_flatten(real_out)),
+ )
+ else:
+ tree_map_(go, fake_out, real_out)
 
  # If a data-dependent op is used in a decomposition, we
  # may need to get the unbacked settings "early"
@@ -1940,13 +1964,15 @@ def go(t: object, real_t: Tensor) -> None:
  )
  ):
  with self:
- return decomposition_table[func](*args, **kwargs)
+ return maybe_propagate_real_tensors(
+ decomposition_table[func](*args, **kwargs)
+ )
 
  with self:
  # Decomposes CompositeImplicitAutograd ops
  r = func.decompose(*args, **kwargs)
  if r is not NotImplemented:
- return r
+ return maybe_propagate_real_tensors(r)
 
  # prims already wrap FakeTensor inputs to FakeTensor outputs
  # and do device logic, we dont need do anything but run them

torch/fx/experimental/symbolic_shapes.py

@@ -2726,6 +2726,7 @@ def _eliminate_unbacked(self, orig_s: sympy.Symbol, new_s: sympy.Expr):
  def set_unbacked_var_to_val(self, k: sympy.Symbol, v: int) -> None:
  """Used only when propagate_real_tensors; registers a value for an
  unbacked symbol, which can be used last resort to resolve hints."""
+ log.info("set_unbacked_var_to_val %s = %s", k, v)
  self.unbacked_var_to_val[k] = sympy.sympify(v)
 
  # Unlike set_replacement, this records a shapeenv event