[functorch] Move cond.py to _cond.py and expose cond() under functorch.experimental.control_flow. (pytorch#89819)

Summary:
Similar to pytorch#88767 we want to reduce the chance that users
accidentally import private functions from `functorch.experimental.cond` as if they were public
interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op.

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

Pull Request resolved: pytorch#89819
Approved by: https://github.com/zou3519

Author: zhxchen17

functorch/experimental/cond.py functorch/experimental/_cond.py

@@ -1,5 +1,3 @@
-# TODO(zhxchen17) Expose API through functorhc.experimental.control_flow
-#                 and rename this file to _cond.py.
 import torch
 
 import torch.utils._pytree as pytree

functorch/experimental/control_flow.py

@@ -1 +1,2 @@
 from ._map import map  # noqa: F401
+from ._cond import cond  # noqa: F401

test/dynamo/test_export.py

@@ -1434,7 +1434,7 @@ def nop(x):
 
     @patch.object(torch._dynamo.config, "capture_scalar_outputs", True)
     def test_export_with_module_layer(self):
-        from functorch.experimental.cond import cond
+        from functorch.experimental.control_flow import cond
 
         def true_fn(layer, val):
             return layer(val) * torch.tensor(2)

test/dynamo/test_misc.py

@@ -2334,7 +2334,7 @@ def f_onnx(x):
         self.assertEqual(f_onnx(input_two_dims), 8)
 
     def test_cond(self):
-        from functorch.experimental.cond import cond
+        from functorch.experimental.control_flow import cond
 
         def true_fn(x):
             return x.sin()
@@ -2352,7 +2352,7 @@ def f(pred, x):
         self.assertTrue(same(torch.sin(torch.tensor([0.25, 0.25])), b))
 
     def test_cond_nested(self):
-        from functorch.experimental.cond import cond
+        from functorch.experimental.control_flow import cond
 
         def true_fn_nested(x):
             return x * 10
@@ -2397,7 +2397,7 @@ def f(pred, pred2, x):
         self.assertTrue(cc.frame_count, 2)
 
     def test_cond_export(self):
-        from functorch.experimental.cond import cond
+        from functorch.experimental.control_flow import cond
 
         def true_fn_nested(x):
             return x * 10
@@ -2442,7 +2442,7 @@ def f(pred, pred2, x):
         )  # * -1 then add x
 
     def test_cond_export_single_arg(self):
-        from functorch.experimental.cond import cond
+        from functorch.experimental.control_flow import cond
 
         def true_fn(x):
             return x

test/functorch/test_control_flow.py

@@ -1,7 +1,7 @@
 # Owner(s): ["module: functorch"]
 import torch
-from functorch.experimental.cond import cond
 from functorch.experimental import control_flow
+from functorch.experimental.control_flow import cond
 from torch.fx.experimental.proxy_tensor import make_fx
 
 from torch.testing._internal.common_utils import run_tests, TestCase