vollo-torch 27.0.0__py3-none-any.whl

vollo_torch/__init__.py

@@ -0,0 +1,105 @@
+from contextvars import ContextVar
+import vollo_torch.fx as fx
+import vollo_torch.nn as nn
+from typing import Sequence
+
+
+class Fp8Weights:
+    """
+    Weights used inside this context will be quantized to an 8-bit format, which may be useful
+    for getting a model to fit on the board.
+
+    Note that this only applies to constant weight matrix multiplications. To ensure predictable
+    behaviour, all MatMuls and Linears that would use dynamic weights are rejected inside this context.
+    """
+
+    _depth: ContextVar[int] = ContextVar("depth", default=0)
+
+    def __enter__(self):
+        self._token = type(self)._depth.set(type(self)._depth.get() + 1)
+        return self
+
+    def __exit__(self, *args):
+        type(self)._depth.reset(self._token)
+
+    @classmethod
+    def is_active(cls) -> bool:
+        return cls._depth.get() > 0
+
+
+class Fp32Activations:
+    """
+    Activations/computations inside this context will be in 32-bit precision. See
+    https://vollo.myrtle.ai/latest/supported-models.html for which operations support this.
+    """
+
+    _depth: ContextVar[int] = ContextVar("depth", default=0)
+
+    def __enter__(self):
+        self._token = type(self)._depth.set(type(self)._depth.get() + 1)
+        return self
+
+    def __exit__(self, *args):
+        type(self)._depth.reset(self._token)
+
+    @classmethod
+    def is_active(cls) -> bool:
+        return cls._depth.get() > 0
+
+
+class CorePartition:
+    """
+    Computations inside this context will be limited to the cores specified in the context's constructor.
+
+    If override=False (default), the CorePartition must specify a subset of the enclosing partition's
+    cores; a ValueError is raised at context entry if this is violated.
+
+    If override=True, the partition overrides the enclosing partition.
+
+    Warning:
+        This class is experimental and is subject to change in future versions.
+    """
+
+    _current_cores: ContextVar[set] = ContextVar("_current_cores", default=None)
+    _current_outer_scan_partition_override: ContextVar[bool] = ContextVar(
+        "_current_outer_scan_partition_override", default=False
+    )
+
+    def __init__(self, cores: Sequence[int], override=False):
+        self.cores = set(cores)
+        if not self.cores:
+            raise ValueError("CorePartition cores must not be empty.")
+        self.override = override
+        self._token = None
+        self._token_outer_scan_partition_override = None
+
+    def __enter__(self):
+        outer_cores = type(self)._current_cores.get()
+        if not self.override and outer_cores is not None:
+            if not self.cores.issubset(outer_cores):
+                raise ValueError(
+                    f"CorePartition cores {self.cores} are not a subset of the enclosing "
+                    f"partition's cores {outer_cores}. "
+                    "Use override=True to override the outer partition."
+                )
+        self._token = type(self)._current_cores.set(self.cores)
+        self._token_outer_scan_partition_override = type(
+            self
+        )._current_outer_scan_partition_override.set(
+            self.override or type(self)._current_outer_scan_partition_override.get()
+        )
+        return self
+
+    def __exit__(self, *args):
+        type(self)._current_cores.reset(self._token)
+        type(self)._current_outer_scan_partition_override.reset(
+            self._token_outer_scan_partition_override
+        )
+
+    @classmethod
+    def current_cores(cls):
+        return cls._current_cores.get()
+
+    @classmethod
+    def current_outer_scan_partition_override(cls):
+        return cls._current_outer_scan_partition_override.get()

vollo_torch/fx/__init__.py

@@ -0,0 +1,49 @@
+from contextvars import ContextVar
+import torch
+import torch.nn as nn
+import torch.fx
+import torch.fx.passes.shape_prop
+
+import vollo_torch.fx.tracer as tracer
+import vollo_torch.fx.util as util
+import vollo_torch.fx.nnir as nnir
+from typing import List, Sequence, Tuple, Union
+
+
+# Use Union in the type hint instead of | because our wheels are for Python 3.7+
+# but | was added in Python 3.10 (Union was added in Python 3.5)
+def prepare_shape(
+    model: nn.Module, *inputs: torch.Tensor
+) -> Tuple[torch.fx.GraphModule, Union[torch.Tensor, Tuple[torch.Tensor, ...]]]:
+    """
+    Trace the model's execution to annotate it with activation shapes. Necessary before converting a model
+    to NNIR.
+    """
+    return _prepare_shape_internal(model, *inputs)
+
+
+def _prepare_shape_internal(
+    model: nn.Module, *inputs: Union[torch.Tensor, Sequence[torch.Tensor]]
+) -> Tuple[
+    torch.fx.GraphModule,
+    Union[torch.Tensor, Tuple[Union[torch.Tensor, Sequence[torch.Tensor]], ...]],
+]:
+    """
+    Same as prepare_shape, but allows more input/output types (such as lists, which aren't allowed as inputs/outputs
+    in normal models, but are allowed as inputs/outputs in a Scan step function)
+    """
+    with torch.no_grad():
+        t = tracer.Tracer()
+        graph = t.trace(model)
+        graph_model = torch.fx.GraphModule(model, graph)
+        outputs = torch.fx.passes.shape_prop.ShapeProp(graph_model).run(*inputs)
+        return (graph_model, outputs)
+
+
+# Re-export
+from vollo_torch.fx.serialization import save, load
+
+is_vollo_tracing: ContextVar[bool] = ContextVar("is_vollo_tracing", default=False)
+"""
+Describes whether the Vollo tracer is currently active. Make sure to use ``.get()`` to read the value.
+"""