Merge branch 'main' into dev/gokul/flmrife-cat-shape

Author: gokulkrishna98

.github/workflows/ci.yml

@@ -0,0 +1,37 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [main]
+  push:
+    branches: [main]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    if: github.repository == 'apple/coreai-torch'
+    runs-on: [self-hosted, macos, tahoe, ARM64]
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+      - name: Ensure uv
+        run: |
+          command -v uv >/dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.local/bin" >> "$GITHUB_PATH"
+      - run: uv run --extra dev ruff check .
+      - run: uv run --extra dev ruff format --check .
+
+  python-test:
+    if: github.repository == 'apple/coreai-torch'
+    runs-on: [self-hosted, macos, tahoe, ARM64]
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+      - name: Ensure uv
+        run: |
+          command -v uv >/dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.local/bin" >> "$GITHUB_PATH"
+      - run: uv run --extra test pytest tests/ -n auto -m "not slow"

coreai_torch/_compression/utils.py

@@ -99,6 +99,7 @@ def repeat_tensor_as(tensor: torch.Tensor, target_shape: torch.Size) -> torch.Te
             )
     return repeated_tensor
 
+
 def wrap_for_parametrization(
     compression_module_class: type[torch.nn.Module],
 ) -> type[torch.nn.Module]:

docs/coreai-core/tutorials/construct-a-graph.ipynb

@@ -76,11 +76,10 @@
     "\n",
     "import numpy as np\n",
     "\n",
-    "from coreai.authoring import AIModelAsset, AIProgram, Module, TensorSpec\n",
-    "\n",
     "# Graph-building primitives — pending re-export from coreai.authoring.\n",
     "from coreai._compiler.dialects import coreai as ops\n",
-    "from coreai._compiler.ir import Value"
+    "from coreai._compiler.ir import Value\n",
+    "from coreai.authoring import AIModelAsset, AIProgram, Module, TensorSpec"
    ]
   },
   {
@@ -175,6 +174,7 @@
     "    ) -> Annotated[Value, output_spec]:\n",
     "        return ops.add(x, x)\n",
     "\n",
+    "\n",
     "module.verify()\n",
     "print(\"Module verified.\")"
    ]

docs/coreai-core/tutorials/run-an-aimodel.ipynb

@@ -64,7 +64,6 @@
     "from pathlib import Path\n",
     "\n",
     "import numpy as np\n",
-    "\n",
     "from coreai.authoring import AIModelAsset\n",
     "from coreai.runtime import InferenceFunction, NDArray\n",
     "\n",
@@ -108,10 +107,10 @@
     "from shutil import rmtree\n",
     "from typing import Annotated\n",
     "\n",
-    "from coreai.authoring import AIProgram, Module, TensorSpec\n",
     "# Pending re-export from coreai.authoring; see the previous tutorial.\n",
     "from coreai._compiler.dialects import coreai as ops\n",
     "from coreai._compiler.ir import Value\n",
+    "from coreai.authoring import AIProgram, Module, TensorSpec\n",
     "\n",
     "# Reconstruct asset.\n",
     "if asset_path.exists():\n",
@@ -126,6 +125,7 @@
     "    ) -> Annotated[Value, TensorSpec(shape=[2, 3], dtype=np.float32, name=\"y\")]:\n",
     "        return ops.add(x, x)\n",
     "\n",
+    "\n",
     "AIProgram(module).save_asset(asset_path)\n",
     "print(f\"created {asset_path}\")"
    ]

docs/getting-started/quickstart.ipynb

@@ -68,8 +68,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
-    "\n",
     "example_input = (torch.randn(1, 10),)\n",
     "exported = torch.export.export(model, args=example_input)"
    ]
@@ -145,10 +143,10 @@
    "outputs": [],
    "source": [
     "import tempfile\n",
-    "import torch\n",
     "from pathlib import Path\n",
     "\n",
     "import numpy as np\n",
+    "import torch\n",
     "from coreai.runtime import NDArray\n",
     "\n",
     "\n",
@@ -240,9 +238,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import tempfile\n",
-    "\n",
-    "\n",
     "async def run():\n",
     "    with tempfile.TemporaryDirectory() as tmpdir:\n",
     "        asset = coreai_program.save_asset(Path(tmpdir) / \"mobilenet_v2_example.aimodel\")\n",
@@ -274,6 +269,7 @@
    "outputs": [],
    "source": [
     "import torch\n",
+    "\n",
     "import coreai_torch\n",
     "\n",
     "model = SimpleModel().eval()\n",

docs/guides/composite-ops.ipynb

@@ -97,11 +97,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "from coreai_torch.composite_ops import RMSNormImpl\n",
-    "\n",
-    "\n",
     "class RMSNorm(nn.Module):\n",
     "    \"\"\"Convenience wrapper that owns the learnable scale parameter.\"\"\"\n",
     "\n",
@@ -149,7 +144,6 @@
    "outputs": [],
    "source": [
     "import torch\n",
-    "import coreai_torch\n",
     "\n",
     "coreai_program = (\n",
     "    TorchConverter()\n",
@@ -185,10 +179,10 @@
    "outputs": [],
    "source": [
     "import tempfile\n",
-    "import torch\n",
     "from pathlib import Path\n",
     "\n",
     "import numpy as np\n",
+    "import torch\n",
     "from coreai.runtime import NDArray\n",
     "\n",
     "\n",

docs/guides/conversion-workflows.ipynb

@@ -86,6 +86,7 @@
    "outputs": [],
    "source": [
     "import torch\n",
+    "\n",
     "import coreai_torch\n",
     "\n",
     "model = MyModel().eval()\n",
@@ -133,7 +134,7 @@
    "source": [
     "import torch\n",
     "import torch.nn as nn\n",
-    "import coreai_torch\n",
+    "\n",
     "from coreai_torch import ExternalizeSpec, TorchConverter\n",
     "from coreai_torch.composite_ops import RMSNormImpl\n",
     "\n",

docs/guides/custom-metal-kernels.ipynb

@@ -122,6 +122,7 @@
    "outputs": [],
    "source": [
     "import torch\n",
+    "\n",
     "from coreai_torch import get_decomp_table\n",
     "\n",
     "model = AddModel().eval()\n",
@@ -182,6 +183,7 @@
    "source": [
     "import torch\n",
     "\n",
+    "\n",
     "def torch_matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n",
     "    return torch.matmul(x, y)\n",
     "\n",
@@ -222,6 +224,7 @@
     "import torch\n",
     "import torch.nn as nn\n",
     "\n",
+    "\n",
     "def torch_sincos(x: torch.Tensor) -> list[torch.Tensor]:\n",
     "    return [torch.sin(x), torch.cos(x)]\n",
     "\n",

docs/guides/custom-op-lowering.ipynb

@@ -92,6 +92,7 @@
    "outputs": [],
    "source": [
     "import torch\n",
+    "\n",
     "from coreai_torch import get_decomp_table\n",
     "\n",
     "model = ScaledAddModel().eval()\n",
@@ -198,9 +199,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import numpy as np\n",
     "import torch\n",
     "import torch.nn as nn\n",
-    "import numpy as np\n",
     "\n",
     "from coreai_torch._utils import get_operand\n",
     "\n",

docs/guides/externalization.ipynb

@@ -86,8 +86,6 @@
     "import torch\n",
     "import torch.nn as nn\n",
     "\n",
-    "from coreai_torch import ExternalizeSpec\n",
-    "\n",
     "\n",
     "class RMSNormComposite(nn.Module):\n",
     "    def __init__(self, axes=-1, eps=1e-5, version=1):\n",
@@ -101,6 +99,7 @@
     "        inv_rms = torch.rsqrt((x_f32 * x_f32).mean(self.axes, keepdim=True) + self.eps)\n",
     "        return (input * inv_rms).to(input.dtype) * scale\n",
     "\n",
+    "\n",
     "model = RMSNormComposite().eval()\n",
     "sample = (torch.randn(10), torch.randn(10))"
    ]