Sync pr-e-python-api with main after PR #254 merge

python/ark/__init__.py

@@ -8,6 +8,7 @@
 
 from .core import version
 from .model import Model
+from .model import set_model, current_model, use_model
 
 __version__ = version()
 

python/ark/data_type.py

@@ -8,6 +8,7 @@
 
 __all__ = [
     "DataType",
+    "bf16",
     "fp16",
     "fp32",
     "int32",

python/ark/model.py

@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from contextlib import contextmanager
 from typing import NewType
 from . import log
 from .core import CoreModel
 
-__all__ = ["Model"]
+__all__ = ["Model", "current_model", "set_model", "use_model"]
 
 ModelState = NewType("ModelState", None)
 
@@ -114,3 +115,45 @@ class ModelState:
     rank: int = 0
     world_size: int = 1
     device_id: int = 0
+
+
+def set_model(model: Model) -> None:
+    """
+    Set the current model.
+
+    Args:
+        model: The model to set as current.
+    """
+    if not isinstance(model, Model):
+        raise log.InvalidUsageError("model must be a Model instance")
+    ModelState.model = model
+
+
+def current_model() -> Model:
+    """
+    Return the current model, creating one if none exists.
+    """
+    return Model.get_model()
+
+
+@contextmanager
+def use_model(model: Model = None):
+    """
+    Context manager that sets *model* as the current model on entry and
+    restores the previous model on exit. If *model* is ``None``, a fresh
+    model is created.
+
+    Args:
+        model: The model to use inside the context. If ``None``, a new
+            ``Model`` is created with the current rank and world size.
+    """
+    prev = ModelState.model
+    if model is None:
+        model = Model(ModelState.rank, ModelState.world_size)
+    elif not isinstance(model, Model):
+        raise log.InvalidUsageError("model must be a Model instance")
+    ModelState.model = model
+    try:
+        yield model
+    finally:
+        ModelState.model = prev

python/ark/ops.py

@@ -49,6 +49,19 @@
 ]
 
 
+def _ensure_ark(t):
+    """
+    If *t* is a ``torch.Tensor``, convert it to an ARK ``Tensor`` via
+    ``Tensor.from_torch()``. Otherwise return *t* unchanged.
+
+    When used on an ``output`` parameter, the returned ARK tensor shares memory
+    with the original torch tensor via ``Tensor.from_torch``.
+    """
+    if not _no_torch and isinstance(t, torch.Tensor):
+        return Tensor.from_torch(t)
+    return t
+
+
 def is_list_or_tuple(obj):
     return isinstance(obj, list) or isinstance(obj, tuple)
 
@@ -60,6 +73,9 @@ def add(
     name: str = "add",
 ) -> Union[Tensor, float]:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if isinstance(input, Tensor) and isinstance(other, Tensor):
         a = input._tensor
         b = other._tensor
@@ -87,6 +103,8 @@ def cast(
     name: str = "cast",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -112,6 +130,8 @@ def copy(
     name: str = "copy",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     if isinstance(input, Tensor):
@@ -126,6 +146,9 @@ def div(
     name: str = "div",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     if isinstance(other, Tensor):
@@ -140,6 +163,9 @@ def embedding(
     name: str = "embedding",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    weight = _ensure_ark(weight)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -153,6 +179,8 @@ def exp(
     name: str = "exp",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().exp(input._tensor, output, name))
@@ -164,6 +192,8 @@ def gelu(
     name: str = "gelu",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().gelu(input._tensor, output, name))
@@ -173,6 +203,8 @@ def identity(
     input: Tensor, deps: List[Tensor] = [], name: str = "identity"
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    deps = [_ensure_ark(d) for d in deps]
     dep_tensors = []
     for dep in deps:
         if not isinstance(dep, Tensor):
@@ -190,6 +222,9 @@ def matmul(
     name: str = "matmul",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -211,6 +246,9 @@ def mul(
     name: str = "mul",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     if isinstance(other, Tensor):
@@ -220,6 +258,7 @@ def mul(
 
 def noop(input: Tensor, name: str = "noop"):
     """ """
+    input = _ensure_ark(input)
     Model.get_model().noop(input._tensor, name)
 
 
@@ -260,6 +299,8 @@ def reduce_max(
     name: str = "reduce_max",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -277,6 +318,8 @@ def reduce_mean(
     name: str = "reduce_mean",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -294,6 +337,8 @@ def reduce_sum(
     name: str = "reduce_sum",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -309,6 +354,8 @@ def relu(
     name: str = "relu",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().relu(input._tensor, output, name))
@@ -333,6 +380,7 @@ def reshape(
     # tensors shape is [128, 64]
     tensor = ark.reshape(tensor, [2, 64, 64])
     """
+    input = _ensure_ark(input)
     if not is_list_or_tuple(shape):
         raise log.InvalidUsageError(
             "shape should be a list or tuple of integers"
@@ -354,6 +402,9 @@ def rope(
     name: str = "rope",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -367,6 +418,8 @@ def rsqrt(
     name: str = "rsqrt",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().rsqrt(input._tensor, output, name))
@@ -376,6 +429,7 @@ def sharding(
     input: Tensor, axis: int, dim_per_shard: int, name: str = "sharding"
 ) -> List[Tensor]:
     """ """
+    input = _ensure_ark(input)
     _tensor_list = Model.get_model().sharding(
         input._tensor, axis, dim_per_shard, name
     )
@@ -388,6 +442,8 @@ def sigmoid(
     name: str = "sigmoid",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().sigmoid(input._tensor, output, name))
@@ -399,6 +455,8 @@ def sqrt(
     name: str = "sqrt",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(Model.get_model().sqrt(input._tensor, output, name))
@@ -411,6 +469,9 @@ def sub(
     name: str = "sub",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    other = _ensure_ark(other)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     if isinstance(other, Tensor):
@@ -442,6 +503,8 @@ def transpose(
     name: str = "transpose",
 ) -> Tensor:
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     if not is_list_or_tuple(perm):
@@ -467,6 +530,8 @@ def send(
     name: str = "send",
 ):
     """ """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     return Tensor(
@@ -479,6 +544,7 @@ def send_done(
     name: str = "send_done",
 ):
     """ """
+    input = _ensure_ark(input)
     return Tensor(Model.get_model().send_done(input._tensor, name))
 
 
@@ -489,6 +555,7 @@ def recv(
     name: str = "recv",
 ):
     """ """
+    output = _ensure_ark(output)
     return Tensor(
         Model.get_model().recv(output._tensor, remote_rank, tag, name)
     )
@@ -591,6 +658,8 @@ def all_reduce(
     Returns:
         Tensor: The reduced tensor.
     """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     _tensor = Model.get_model().all_reduce(
@@ -625,6 +694,8 @@ def all_reduce_packet(
     Returns:
         Tensor: The reduced tensor.
     """
+    input = _ensure_ark(input)
+    output = _ensure_ark(output)
     if output is not NullTensor:
         output = output._tensor
     _tensor = Model.get_model().all_reduce_packet(

python/ark/planner.py

@@ -232,8 +232,10 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
 
 class Planner(CorePlanner):
-    def __init__(self, device_id: int = 0):
-        compressed = Model.get_model().compress()
+    def __init__(self, device_id: int = 0, model: Model = None):
+        if model is None:
+            model = Model.get_model()
+        compressed = model.compress()
         super().__init__(compressed, device_id)
 
     def install_config_rule(self, rule: Callable[[str, str], str]):