Use safe_is_leaf to test leafness

This fixes one of the problems in #101160 (comment)
but I don't have a test case because the full example is fairly
difficult to minify.

Signed-off-by: Edward Z. Yang <ezyangmeta.com>

ghstack-source-id: 6a4e9aa7bb7693d87971caecab7c9c72e9accba8
Pull Request resolved: #102706

torch/_functorch/aot_autograd.py

@@ -19,6 +19,7 @@
 import torch.utils._pytree as pytree
 import torch.utils.dlpack
 from torch import Tensor
+from torch._subclasses.meta_utils import safe_is_leaf
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.utils import dynamo_timed, lazy_format_graph_code
 from torch._guards import detect_fake_mode, tracing
@@ -748,7 +749,7 @@ def inner(*flat_args):
                 mutates_metadata = False
 
             input_info.append(InputAliasInfo(
-                is_leaf=isinstance(arg, torch.Tensor) and arg.is_leaf,
+                is_leaf=isinstance(arg, torch.Tensor) and safe_is_leaf(arg),
                 mutates_data=mutates_data,
                 mutates_metadata=mutates_metadata
             ))