Repository: pytorch
Issue: [export] non-strict export error on torch.distributions.Normal
Labels: triaged,oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

The following issues showed up when I ran **non-strict** export on models `soft_actor_critic` and `drq` in `torchbench`. There is a data-dependent expression when calling `torch.distributions.Normal`. However, running strict mode doesn't have such issues

**Repro**
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import distributions as pyd

class SquashedNormal(pyd.transformed_distribution.TransformedDistribution):
    def __init__(self, loc, scale, tanh_transform_clamp=(-0.99, 0.99)):
        self.loc = loc
        self.scale = scale
        self.tanh_transform_clamp = tanh_transform_clamp
        self.base_dist = pyd.Normal(loc, scale)
        super().__init__(self.base_dist, [])

    @property
    def mean(self):
        mu = self.loc
        for tr in self.transforms:
            mu = tr(mu)
        return mu


def _squashed_normal_flatten(t: SquashedNormal):
    return [t.loc, t.scale], t.tanh_transform_clamp


def _squashed_normal_unflatten(values, context):
    return SquashedNormal(*values, context)


torch.utils._pytree.register_pytree_node(
    SquashedNormal,
    _squashed_normal_flatten,
    _squashed_normal_unflatten,
    serialized_type_name=f"{SquashedNormal.__module__}.{SquashedNormal.__name__}",
)

class StochasticActor(nn.Module):
    def __init__(
        self,
        state_space_size=4,
        act_space_size=2,
    ):
        super().__init__()
        self.fc = nn.Linear(state_space_size, act_space_size)

    def forward(self, state):
        out = F.relu(self.fc(state))
        mu, log_std = out.chunk(2, dim=1)
        log_std = torch.tanh(log_std)
        std = (-10 + 6*(log_std+1)).exp()
        dist = SquashedNormal(mu, std)
        return dist


model = StochasticActor()
ep = torch.export.export(model, (torch.randn(1,4),), strict=False)
print (ep)

```
**Error Msg**
```
W0903 15:10:15.901000 2821440 torch/fx/experimental/symbolic_shapes.py:5128] failed during evaluate_expr(Eq(u0, 1), hint=None, expect_rational=True, size_oblivious=False, forcing_spec=False
E0903 15:10:15.901000 2821440 torch/fx/experimental/recording.py:298] failed while running evaluate_expr(*(Eq(u0, 1), None), **{'fx_node': False})
Traceback (most recent call last):
  File "/home/yimingzhou/pytorch/torch/_export/non_strict_utils.py", line 513, in __torch_function__
    return func(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/fx/experimental/sym_node.py", line 451, in guard_bool
    r = self.shape_env.evaluate_expr(self.expr, self.hint, fx_node=self.fx_node)
  File "/home/yimingzhou/pytorch/torch/fx/experimental/recording.py", line 262, in wrapper
    return retlog(fn(*args, **kwargs))
  File "/home/yimingzhou/pytorch/torch/fx/experimental/symbolic_shapes.py", line 5126, in evaluate_expr
    return self._evaluate_expr(orig_expr, hint, fx_node, expect_rational, size_oblivious, forcing_spec=forcing_spec)
  File "/home/yimingzhou/pytorch/torch/fx/experimental/symbolic_shapes.py", line 5244, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1)).  (Size-like symbols: none)

Potential framework code culprit (scroll up for full backtrace):
  File "/home/yimingzhou/pytorch/torch/_export/non_strict_utils.py", line 513, in __torch_function__
    return func(*args, **kwargs)

For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="u0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The following call raised this error:
  File "/home/yimingzhou/pytorch/pyd_normal_repro.py", line 11, in __init__
    self.base_dist = pyd.Normal(loc, scale)


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/yimingzhou/pytorch/pyd_normal_repro.py", line 57, in <module>
    ep = torch.export.export(model, (torch.randn(1,4),), strict=False)
  File "/home/yimingzhou/pytorch/torch/export/__init__.py", line 258, in export
    return _export(
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 1007, in wrapper
    raise e
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 980, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/export/exported_program.py", line 105, in wrapper
    return fn(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 1915, in _export
    export_artifact = export_func(  # type: ignore[operator]
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 1719, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 626, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 1648, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
  File "/home/yimingzhou/pytorch/torch/_functorch/aot_autograd.py", line 1246, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/home/yimingzhou/pytorch/torch/_functorch/aot_autograd.py", line 1480, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/home/yimingzhou/pytorch/torch/_functorch/aot_autograd.py", line 522, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/home/yimingzhou/pytorch/torch/_functorch/aot_autograd.py", line 623, in _create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
  File "/home/yimingzhou/pytorch/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 168, in inner
    flat_f_outs = f(*flat_f_args)
  File "/home/yimingzhou/pytorch/torch/_functorch/_aot_autograd/utils.py", line 182, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 863, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/home/yimingzhou/pytorch/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/export/_trace.py", line 1635, in forward
    tree_out = self._export_root(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/yimingzhou/pytorch/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/yimingzhou/pytorch/pyd_normal_repro.py", line 52, in forward
    dist = SquashedNormal(mu, std)
  File "/home/yimingzhou/pytorch/pyd_normal_repro.py", line 11, in __init__
    self.base_dist = pyd.Normal(loc, scale)
  File "/home/yimingzhou/pytorch/torch/distributions/normal.py", line 59, in __init__
    super().__init__(batch_shape, validate_args=validate_args)
  File "/home/yimingzhou/pytorch/torch/distributions/distribution.py", line 70, in __init__
    if not valid.all():
  File "/home/yimingzhou/pytorch/torch/_export/non_strict_utils.py", line 515, in __torch_function__
    _suggest_fixes_for_data_dependent_error_non_strict(e)
  File "/home/yimingzhou/pytorch/torch/fx/experimental/symbolic_shapes.py", line 5598, in _suggest_fixes_for_data_dependent_error_non_strict
    for path, leaf in pytree.tree_leaves_with_path(val):
  File "/home/yimingzhou/pytorch/torch/utils/_pytree.py", line 1544, in tree_leaves_with_path
    return list(_generate_key_paths((), tree, is_leaf))
  File "/home/yimingzhou/pytorch/torch/utils/_pytree.py", line 1570, in _generate_key_paths
    raise ValueError(
ValueError: Did not find a flatten_with_keys_fn for type: <class '__main__.SquashedNormal'>. Please pass a flatten_with_keys_fn argument to register_pytree_node.
```

### Versions

main

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4