Patch executorch.pytree dependency

Summary: For running the AOT side of executorch w/o needing to build the pybindings, we can load the executorch.pytree.pybinding functions as the torch.utils._pytree functions.

Reviewed By: larryliu0820

Differential Revision: D47136847

fbshipit-source-id: 0da9ce09636f91ed3368c83f1867f6c70544d736

Author: angelayi

pytree/__init__.py

@@ -1,13 +1,43 @@
 # flake8: noqa: F401
 
-# pyre-fixme[21]: Could not find module `executorch.pytree.pybindings`.
-# @manual=//executorch/pytree:pybindings
-from executorch.pytree.pybindings import (
-    broadcast_to_and_flatten as broadcast_to_and_flatten,
-    from_str as from_str,
-    register_custom as register_custom,
-    tree_flatten as tree_flatten,
-    tree_map as tree_map,
-    tree_unflatten as tree_unflatten,
-    TreeSpec as TreeSpec,
-)
+import warnings
+
+try:
+    """
+    Internally we link the respective c++ library functions but for the OSS pip
+    build we will just use the python library for now. The python library is not
+    exactly the same so it will not work for the runtime, but it'll be fine for
+    now as in most cases the runtime will not need it.
+    """
+
+    # pyre-fixme[21]: Could not find module `executorch.pytree.pybindings`.
+    # @manual=//executorch/pytree:pybindings
+    from executorch.pytree.pybindings import (
+        broadcast_to_and_flatten as broadcast_to_and_flatten,
+        from_str as from_str,
+        register_custom as register_custom,
+        tree_flatten as tree_flatten,
+        tree_map as tree_map,
+        tree_unflatten as tree_unflatten,
+        TreeSpec as TreeSpec,
+    )
+except:
+    warnings.warn(
+        "Unable to import executorch.pytree, using native torch pytree instead."
+    )
+
+    from torch.utils._pytree import (
+        _broadcast_to_and_flatten,
+        _register_pytree_node,
+        pytree_to_str,
+        str_to_pytree,
+        tree_flatten,
+        tree_map,
+        tree_unflatten,
+        TreeSpec,
+    )
+
+    broadcast_to_and_flatten = _broadcast_to_and_flatten
+    from_str = str_to_pytree
+    register_custom = _register_pytree_node
+    TreeSpec.to_str = pytree_to_str  # pyre-ignore