[monarch] initial commit of tensor_engine submodule

Pull Request resolved: #458

Now that the actor layer is place, time to refactor the tensor engine following the same principles:

- Every symbol has at most two locations: a private location (which should be exactly the same as its location in the source tree), and optionally a public location in the user-facing API.
- We discourage end-users from depending on private locations via underscores, per python convention.

In particular, we violate principle 2 in many places in the tensor engine—all of our top-level files/modules are visible if you do `from monarch.foo import bar`, and we are sloppy about the distinction of public and private.

So this refactor will be a little bit painful. I will approach it in parts. This diff defines the `tensor_engine` submodule and moves `common` to `monarch.tensor_engine._common`.

Differential Revision: [D77847587](https://our.internmc.facebook.com/intern/diff/D77847587/)
ghstack-source-id: 294872880

Author: suo

monarch_extension/src/convert.rs

@@ -44,7 +44,7 @@ struct MessageParser<'a> {
 fn create_function(obj: Bound<'_, PyAny>) -> PyResult<ResolvableFunction> {
     let cloudpickle = obj
         .py()
-        .import("monarch.common.function")?
+        .import("monarch._src.tensor_engine.common.function")?
         .getattr("ResolvableFromCloudpickle")?;
     if obj.is_instance(&cloudpickle)? {
         Ok(ResolvableFunction::Cloudpickle(Cloudpickle::new(
@@ -102,7 +102,7 @@ impl<'a> MessageParser<'a> {
         let referenceable = self
             .current
             .py()
-            .import("monarch.common.reference")?
+            .import("monarch._src.tensor_engine.common.reference")?
             .getattr("Referenceable")?;
         let mut flat: Vec<Option<Ref>> = vec![];
         for x in output_tuple.0.try_iter()? {
@@ -198,8 +198,8 @@ static CONVERT_MAP: OnceLock<HashMap<u64, FnType>> = OnceLock::new();
 
 fn create_map(py: Python) -> HashMap<u64, FnType> {
     let messages = py
-        .import("monarch.common.messages")
-        .expect("import monarch.common.messages");
+        .import("monarch._src.tensor_engine.common.messages")
+        .expect("import monarch._src.tensor_engine.common.messages");
     let mut m: HashMap<u64, FnType> = HashMap::new();
     let key = |name: &str| {
         messages

python/monarch/__init__.py

@@ -34,9 +34,9 @@
     from monarch import timer
     from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator
     from monarch._src.actor.shape import NDSlice, Shape
-    from monarch.common._coalescing import coalescing
+    from monarch._src.tensor_engine.common._coalescing import coalescing
 
-    from monarch.common.device_mesh import (
+    from monarch._src.tensor_engine.common.device_mesh import (
         DeviceMesh,
         get_active_mesh,
         no_mesh,
@@ -45,17 +45,23 @@
         to_mesh,
     )
 
-    from monarch.common.function import resolvers as function_resolvers
+    from monarch._src.tensor_engine.common.function import (
+        resolvers as function_resolvers,
+    )
 
-    from monarch.common.future import Future
+    from monarch._src.tensor_engine.common.future import Future
 
-    from monarch.common.invocation import RemoteException
-    from monarch.common.opaque_ref import OpaqueRef
-    from monarch.common.pipe import create_pipe, Pipe, remote_generator
-    from monarch.common.remote import remote
-    from monarch.common.selection import Selection
-    from monarch.common.stream import get_active_stream, Stream
-    from monarch.common.tensor import reduce, reduce_, Tensor
+    from monarch._src.tensor_engine.common.invocation import RemoteException
+    from monarch._src.tensor_engine.common.opaque_ref import OpaqueRef
+    from monarch._src.tensor_engine.common.pipe import (
+        create_pipe,
+        Pipe,
+        remote_generator,
+    )
+    from monarch._src.tensor_engine.common.remote import remote
+    from monarch._src.tensor_engine.common.selection import Selection
+    from monarch._src.tensor_engine.common.stream import get_active_stream, Stream
+    from monarch._src.tensor_engine.common.tensor import reduce, reduce_, Tensor
     from monarch.fetch import fetch_shard, inspect, show
     from monarch.gradient_generator import grad_function, grad_generator
     from monarch.notebook import mast_mesh, reserve_torchx as mast_reserve
@@ -72,29 +78,41 @@
 
 
 _public_api = {
-    "coalescing": ("monarch.common._coalescing", "coalescing"),
-    "remote": ("monarch.common.remote", "remote"),
-    "DeviceMesh": ("monarch.common.device_mesh", "DeviceMesh"),
-    "get_active_mesh": ("monarch.common.device_mesh", "get_active_mesh"),
-    "no_mesh": ("monarch.common.device_mesh", "no_mesh"),
-    "RemoteProcessGroup": ("monarch.common.device_mesh", "RemoteProcessGroup"),
-    "function_resolvers": ("monarch.common.function", "resolvers"),
-    "Future": ("monarch.common.future", "Future"),
-    "RemoteException": ("monarch.common.invocation", "RemoteException"),
+    "coalescing": ("monarch._src.tensor_engine.common._coalescing", "coalescing"),
+    "remote": ("monarch._src.tensor_engine.common.remote", "remote"),
+    "DeviceMesh": ("monarch._src.tensor_engine.common.device_mesh", "DeviceMesh"),
+    "get_active_mesh": (
+        "monarch._src.tensor_engine.common.device_mesh",
+        "get_active_mesh",
+    ),
+    "no_mesh": ("monarch._src.tensor_engine.common.device_mesh", "no_mesh"),
+    "RemoteProcessGroup": (
+        "monarch._src.tensor_engine.common.device_mesh",
+        "RemoteProcessGroup",
+    ),
+    "function_resolvers": ("monarch._src.tensor_engine.common.function", "resolvers"),
+    "Future": ("monarch._src.tensor_engine.common.future", "Future"),
+    "RemoteException": (
+        "monarch._src.tensor_engine.common.invocation",
+        "RemoteException",
+    ),
     "Shape": ("monarch._src.actor.shape", "Shape"),
     "NDSlice": ("monarch._src.actor.shape", "NDSlice"),
-    "Selection": ("monarch.common.selection", "Selection"),
-    "OpaqueRef": ("monarch.common.opaque_ref", "OpaqueRef"),
-    "create_pipe": ("monarch.common.pipe", "create_pipe"),
-    "Pipe": ("monarch.common.pipe", "Pipe"),
-    "remote_generator": ("monarch.common.pipe", "remote_generator"),
-    "get_active_stream": ("monarch.common.stream", "get_active_stream"),
-    "Stream": ("monarch.common.stream", "Stream"),
-    "Tensor": ("monarch.common.tensor", "Tensor"),
-    "reduce": ("monarch.common.tensor", "reduce"),
-    "reduce_": ("monarch.common.tensor", "reduce_"),
-    "to_mesh": ("monarch.common.device_mesh", "to_mesh"),
-    "slice_mesh": ("monarch.common.device_mesh", "slice_mesh"),
+    "Selection": ("monarch._src.tensor_engine.common.selection", "Selection"),
+    "OpaqueRef": ("monarch._src.tensor_engine.common.opaque_ref", "OpaqueRef"),
+    "create_pipe": ("monarch._src.tensor_engine.common.pipe", "create_pipe"),
+    "Pipe": ("monarch._src.tensor_engine.common.pipe", "Pipe"),
+    "remote_generator": ("monarch._src.tensor_engine.common.pipe", "remote_generator"),
+    "get_active_stream": (
+        "monarch._src.tensor_engine.common.stream",
+        "get_active_stream",
+    ),
+    "Stream": ("monarch._src.tensor_engine.common.stream", "Stream"),
+    "Tensor": ("monarch._src.tensor_engine.common.tensor", "Tensor"),
+    "reduce": ("monarch._src.tensor_engine.common.tensor", "reduce"),
+    "reduce_": ("monarch._src.tensor_engine.common.tensor", "reduce_"),
+    "to_mesh": ("monarch._src.tensor_engine.common.device_mesh", "to_mesh"),
+    "slice_mesh": ("monarch._src.tensor_engine.common.device_mesh", "slice_mesh"),
     "call_on_shard_and_fetch": ("monarch.fetch", "call_on_shard_and_fetch"),
     "fetch_shard": ("monarch.fetch", "fetch_shard"),
     "inspect": ("monarch.fetch", "inspect"),

python/monarch/_rust_bindings/monarch_extension/client.pyi

@@ -8,7 +8,7 @@ from typing import Any, ClassVar, Dict, final, List, NamedTuple, Union
 
 from monarch._rust_bindings.monarch_extension.tensor_worker import Ref
 from monarch._rust_bindings.monarch_messages.debugger import DebuggerActionType
-from monarch._src.actor._extension.monarch_hyperactor.proc import (
+from monarch._src.actor._extension.monarch_hyperactor.proc import (  # @manual=//monarch/python/monarch/_src/actor:actor
     ActorId,
     Proc,
     Serialized,

python/monarch/common/__init__.py renamed to python/monarch/_src/tensor_engine/__init__.py

@@ -24,22 +24,17 @@
 )
 
 import torch
-from monarch.common import messages
 
-from monarch.common.fake import fake_call
-from monarch.common.function_caching import (
-    hashable_tensor_flatten,
-    TensorGroup,
-    TensorGroupPattern,
-)
-from monarch.common.tensor import InputChecker, Tensor
-from monarch.common.tree import flatten
+from . import messages
 
-if TYPE_CHECKING:
-    from monarch.common.client import Recorder
-    from monarch.common.recording import Recording
+from .fake import fake_call
+from .function_caching import hashable_tensor_flatten, TensorGroup, TensorGroupPattern
+from .tensor import InputChecker, Tensor
+from .tree import flatten
 
-    from .client import Client
+if TYPE_CHECKING:
+    from .client import Client, Recorder
+    from .recording import Recording
 
 _coalescing = None