Merge branch 'main' into dev/gokul/smolvlm-masked-scatter

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/_aten_to_core.py

@@ -644,17 +644,20 @@ def replace_arange_start_step(
         else coreai.constant(1, dtype=start.type.element_type)
     )
 
-    # Squeeze rank-1 inputs to 0D scalars for coreai.range_.
+    # coreai.range_ requires scalar (rank-0) operands that share an element
+    # type. aten.arange promotes mixed-type scalars internally; we replicate
+    # that here by squeezing each operand to rank-0 and casting to the FX
+    # node's output dtype before the op.
+    target_type = get_output_element_type_from_node(node)
+
     def to_scalar(v: Value) -> Value:
         if v.type.rank > 0:
-            return coreai.shrink_dims(v, list(range(v.type.rank)))
+            v = coreai.shrink_dims(v, list(range(v.type.rank)))
+        if v.type.element_type != target_type:
+            v = coreai.cast(v, target_type)
         return v
 
-    result = coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
-    target_type = get_output_element_type_from_node(node)
-    if result.type.element_type != target_type:
-        result = coreai.cast(result, target_type)
-    return result
+    return coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
 
 
 def replace_batch_norm(
@@ -814,6 +817,7 @@ def replace_binary_ops(
         "pow.Scalar": coreai.broadcasting_pow,
         "pow.Tensor_Tensor": coreai.broadcasting_pow,
         "pow.Tensor_Scalar": coreai.broadcasting_pow,
+        "pow": coreai.broadcasting_pow,
         "sub.Tensor": coreai.broadcasting_sub,
         "sub.Scalar": coreai.broadcasting_sub,
         "sub": coreai.broadcasting_sub,
@@ -944,6 +948,55 @@ def replace_cat(values_map: dict[str, Value], node: fx.Node, loc: Location) -> V
 
     rank = inputs[0].type.rank
     dim = dim + rank if dim < 0 else dim
+
+    # coreai.concat requires all non-concat dims to be provably equal across
+    # inputs. Under dynamic shapes, one branch can carry a dynamic non-concat
+    # axis while a sibling has a static size for the same axis — the dynamic
+    # side must in fact equal that static size, but the type system doesn't
+    # know it. Reshape such inputs to the known static size before the concat.
+    # Multiple distinct static sizes on one axis is a real mismatch and is
+    # left for the dialect verifier to reject.
+    dyn = ShapedType.get_dynamic_size()
+
+    def known_static(axis: int) -> int | None:
+        if axis == dim:
+            return None
+        sizes = {inp.type.shape[axis] for inp in inputs if inp.type.shape[axis] != dyn}
+        return next(iter(sizes)) if len(sizes) == 1 else None
+
+    statics = [known_static(a) for a in range(rank)]
+    promoted: list[Value] = []
+    for inp in inputs:
+        new_shape = [
+            statics[a]
+            if statics[a] is not None and inp.type.shape[a] == dyn
+            else inp.type.shape[a]
+            for a in range(rank)
+        ]
+        if new_shape != list(inp.type.shape):
+            if all(s != dyn for s in new_shape):
+                # All axes static post-promotion: list-form reshape packs
+                # the shape into an int32 constant tensor.
+                inp = coreai.reshape(inp, new_shape)
+            else:
+                # Mixed static / dynamic post-promotion: build the shape
+                # vector at runtime by mixing the input's actual sizes
+                # (via coreai.get_shape) for the still-dynamic axes with
+                # constants for the promoted axes.
+                runtime_shape = coreai.cast(coreai.get_shape(inp), dtype=np.int32)
+                parts = [
+                    coreai.constant([s], dtype=np.int32)
+                    if s != dyn
+                    else coreai.slice_(runtime_shape, [a], [a + 1], [1])
+                    for a, s in enumerate(new_shape)
+                ]
+                result_type = RankedTensorType.get(new_shape, inp.type.element_type)
+                inp = coreai.ReshapeOp(
+                    inp, coreai.concat(0, parts), results=[result_type]
+                ).result
+        promoted.append(inp)
+    inputs = promoted
+
     return coreai.concat(dim, inputs)
 
 
@@ -2232,11 +2285,29 @@ def replace_remainder(
 
 def replace_repeat(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
     x = _get_operand(values_map, node, 0)
-    repeats = np.array(node.args[1], dtype=np.uint32)
-    extra_dims = len(repeats) - x.type.rank
+    repeat_args = list(node.args[1])
+    extra_dims = len(repeat_args) - x.type.rank
     if extra_dims > 0:
         x = coreai.expand_dims(x, list(range(extra_dims)))
-    return coreai.tile(x, repeats)
+
+    if all(isinstance(r, int) for r in repeat_args):
+        return coreai.tile(x, np.array(repeat_args, dtype=np.uint32))
+
+    # At least one repeat is a SymInt fx.Node — build a rank-1 uint32 dim
+    # vector at runtime, with per-axis constants for plain ints and the
+    # resolved Value (cast to uint32, lifted to rank-1 if scalar) for
+    # SymInts. coreai.tile accepts a runtime Value for its dims.
+    chunks: list[Value] = []
+    for r in repeat_args:
+        if isinstance(r, int):
+            chunks.append(coreai.constant([r], dtype=np.uint32))
+        else:
+            assert isinstance(r, fx.Node)
+            v = coreai.cast(values_map[r.name], dtype=np.uint32)
+            if v.type.rank == 0:
+                v = coreai.reshape(v, [1])
+            chunks.append(v)
+    return coreai.tile(x, coreai.concat(0, chunks))
 
 
 def replace_round_decimals(
@@ -2635,6 +2706,7 @@ def replace_unary_ops(
         "log.default": coreai.log,
         "relu.default": coreai.relu,
         "round.default": coreai.round_,
+        "round": coreai.round_,
         "rsqrt.default": coreai.rsqrt,
         "sigmoid.default": coreai.sigmoid,
         "silu.default": coreai.silu,
@@ -3475,13 +3547,15 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "pow.Scalar": replace_binary_ops,
     "pow.Tensor_Scalar": replace_binary_ops,
     "pow.Tensor_Tensor": replace_binary_ops,
+    "pow": replace_binary_ops,
     "prod.default": replace_prod_default,
     "prod.dim_int": replace_prod_dim_int,
     "reciprocal.default": replace_reciprocal,
     "relu.default": replace_unary_ops,
     "remainder.Tensor": replace_remainder,
     "round.default": replace_unary_ops,
     "round.decimals": replace_round_decimals,
+    "round": replace_unary_ops,
     "repeat.default": replace_repeat,
     "rsqrt.default": replace_unary_ops,
     "scaled_dot_product_attention.default": replace_sdpa,

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]:

coreai_torch/_utils.py

@@ -175,6 +175,23 @@ def __exit__(self, *exc: Any) -> None:
         self.close()
 
 
+def to_rank1_int32(v: Value) -> Value:
+    """Coerce a SymInt-derived Value to canonical rank-1 si32 form.
+
+    Dim-vector concats (used to build shape operands for ``coreai.reshape``,
+    ``coreai.interpolate``, etc.) require all inputs to share rank and
+    element type. SymInt values can arrive rank-0 (e.g. from
+    ``aten._local_scalar_dense``) or with a different int variant. This
+    helper produces the form that aligns with ``coreai.constant([i],
+    dtype=np.int32)`` and ``replace_sym_size_int``.
+    """
+    if v.type.rank == 0:
+        v = coreai.reshape(v, [1])
+    if v.type.element_type != IntegerType.get_signed(32):
+        v = coreai.cast(v, np.int32)
+    return v
+
+
 def upsample_build_output_shape_dynamic(
     x: Value, out_h: int | Value, out_w: int | Value
 ) -> Value:
@@ -192,8 +209,8 @@ def upsample_build_output_shape_dynamic(
     )
     shape = coreai.cast(coreai.get_shape(x), dtype=np.int32)
     non_spatial = coreai.slice_(shape, [0], [2], [1])
-    h = [out_h] if isinstance(out_h, int) else out_h
-    w = [out_w] if isinstance(out_w, int) else out_w
+    h = [out_h] if isinstance(out_h, int) else to_rank1_int32(out_h)
+    w = [out_w] if isinstance(out_w, int) else to_rank1_int32(out_w)
     return coreai.concat(0, [non_spatial, h, w])
 
 
@@ -986,9 +1003,13 @@ def get_operand(
     if isinstance(arg, fx.Node):
         return values_map[arg.name]
     if isinstance(arg, list) and any(isinstance(e, fx.Node) for e in arg):
-        # Mixed list: resolve fx.Node elements via values_map, keep ints as constants.
+        # Mixed list: SymInt fx.Nodes + plain ints. Concat the two sources
+        # into a single rank-1 si32 dim vector. Both branches must produce
+        # the same canonical form so the concat verifier accepts them.
         dim_vals = [
-            values_map[e.name] if isinstance(e, fx.Node) else coreai.constant([e])
+            to_rank1_int32(values_map[e.name])
+            if isinstance(e, fx.Node)
+            else coreai.constant([e], dtype=np.int32)
             for e in arg
         ]
         return coreai.concat(0, dim_vals) if len(dim_vals) > 1 else dim_vals[0]

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",