[pytree] Only import optree if it's used (#131478)

torch.utils._pytree imports optree if it's available. Instead, we change
it to if it gets used. The motivation for this is better isolation.

Test Plan:
- new tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131478
Approved by: https://github.com/albanD

torch/utils/_cxx_pytree.py

@@ -32,8 +32,9 @@ from typing_extensions import deprecated
 import optree
 from optree import PyTreeSpec  # direct import for type annotations
 
-from torch.utils._pytree import KeyEntry
+import torch.utils._pytree as _pytree
 
+from torch.utils._pytree import KeyEntry
 
 __all__ = [
     "PyTree",
@@ -999,3 +1000,8 @@ def keystr(kp: KeyPath) -> str:
 def key_get(obj: Any, kp: KeyPath) -> Any:
     """Given an object and a key path, return the value at the key path."""
     raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.")
+
+
+_pytree._cxx_pytree_imported = True
+for args, kwargs in _pytree._cxx_pytree_pending_imports:
+    _private_register_pytree_node(*args, **kwargs)