transpiler: op routing, fusion & JAX-capture correctness

python/cactus/transpile/aten_ops.py

@@ -123,6 +123,8 @@ def canonical_torch_op(target: Any) -> str:
     AtenOpRule("aten.mean", "mean"),
     AtenOpRule("aten.all", "min"),
     AtenOpRule("aten.var", "variance"),
+    AtenOpRule("aten.minimum", "minimum"),
+    AtenOpRule("aten.maximum", "maximum"),
     AtenOpRule("aten.min", "min"),
     AtenOpRule("aten.amin", "min"),
     AtenOpRule("aten.max", "max"),
@@ -131,6 +133,8 @@ def canonical_torch_op(target: Any) -> str:
     AtenOpRule("aten.split_with_sizes", "split_with_sizes"),
     AtenOpRule("aten.chunk", "chunk"),
     AtenOpRule("aten.ones", "ones"),
+    AtenOpRule("aten.slice_scatter", "slice_scatter"),
+    AtenOpRule("aten.select_scatter", "select_scatter"),
     AtenOpRule("aten.slice", "slice"),
     AtenOpRule("aten.select", "index"),
     AtenOpRule("aten.gather", "gather"),
@@ -151,6 +155,11 @@ def canonical_torch_op(target: Any) -> str:
 )
 
 
+_PREFIX_RULES_BY_LENGTH: tuple[AtenOpRule, ...] = tuple(
+    sorted(_PREFIX_RULES, key=lambda rule: len(rule.torch_prefix), reverse=True)
+)
+
+
 def canonical_ir_op(target: Any) -> str:
     torch_op = canonical_torch_op(target)
     exact = _EXACT_RULES.get(torch_op)
@@ -165,7 +174,7 @@ def canonical_ir_op(target: Any) -> str:
             "aten.lt": "less",
             "aten.le": "less_equal",
         }[torch_op]
-    for rule in _PREFIX_RULES:
+    for rule in _PREFIX_RULES_BY_LENGTH:
         if torch_op.startswith(rule.torch_prefix):
             return rule.ir_op
     return torch_op

python/cactus/transpile/capture_jax.py

@@ -1358,7 +1358,8 @@ def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode |
         limits = tuple(int(value) for value in params["limit_indices"])
         raw_strides = params.get("strides")
         strides = (1,) * len(starts) if raw_strides is None else tuple(int(value) for value in raw_strides)
-        changed_axes = [axis for axis, (start, limit) in enumerate(zip(starts, limits, strict=True)) if start != 0 or limit != getattr(eqn.invars[0].aval, "shape", ())[axis]]
+        input_shape = getattr(eqn.invars[0].aval, "shape", ())
+        changed_axes = [axis for axis, (start, limit, stride) in enumerate(zip(starts, limits, strides, strict=True)) if start != 0 or limit != input_shape[axis] or stride != 1]
         if not changed_axes:
             _alias_output(ctx, eqn.outvars[0], inputs[0])
             return None
@@ -1400,18 +1401,19 @@ def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode |
         if index_shape and index_shape[-1] == 1:
             squeezed_id = f"{indices_id}__squeezed"
             squeeze_shape = tuple(int(dim) for dim in index_shape[:-1])
-            squeeze_node = IRNode(
-                node_id,
-                "reshape",
-                [indices_id],
-                [squeezed_id],
-                attrs={"shape": squeeze_shape},
-                meta=meta,
-            )
-            _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,))
-            graph.values[squeezed_id].shape = squeeze_shape
-            if indices_id in ctx.gather_index_aliases:
-                ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id]
+            if squeezed_id not in graph.values:
+                squeeze_node = IRNode(
+                    node_id,
+                    "reshape",
+                    [indices_id],
+                    [squeezed_id],
+                    attrs={"shape": squeeze_shape},
+                    meta=meta,
+                )
+                _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,))
+                graph.values[squeezed_id].shape = squeeze_shape
+                if indices_id in ctx.gather_index_aliases:
+                    ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id]
             indices_id = squeezed_id
             node_id = ctx.node_id("gather_embedding")
         indices_id = ctx.gather_index_aliases.get(indices_id, indices_id)

python/cactus/transpile/fusion/rms_norm.py

@@ -107,8 +107,27 @@ def _extract_rsqrt_chain(graph: IRGraph, value_id: str) -> dict[str, object] | N
     pow2_node = producer(graph, mean_node.inputs[0])
     if pow2_node is None or pow2_node.op != "pow" or float(pow2_node.attrs.get("exponent", 0.0)) != 2.0:
         return None
+    if not _reduces_last_dim(graph, mean_node, pow2_node.inputs[0]):
+        return None
     return {
         "pow_input_value_id": pow2_node.inputs[0],
         "eps": float(add_node.attrs.get("value", 0.0)),
         "node_ids": (pow_node.id, add_node.id, mean_node.id, pow2_node.id),
     }
+
+
+def _reduces_last_dim(graph: IRGraph, mean_node: IRNode, input_value_id: str) -> bool:
+    value = graph.values.get(input_value_id)
+    if value is None or value.shape is None:
+        return False
+    rank = len(value.shape)
+    if rank == 0:
+        return False
+    axis = mean_node.attrs.get("axis")
+    if isinstance(axis, (list, tuple)):
+        if len(axis) != 1:
+            return False
+        axis = axis[0]
+    if not isinstance(axis, int):
+        return False
+    return axis % rank == rank - 1

python/cactus/transpile/importers.py

@@ -986,6 +986,11 @@ def import_linear(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[in
 
 
 def import_addmm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None:
+    kwargs = getattr(node, "kwargs", {})
+    beta = _extract_numeric_literal(node.args[3]) if len(node.args) > 3 else _extract_numeric_literal(kwargs.get("beta"))
+    alpha = _extract_numeric_literal(node.args[4]) if len(node.args) > 4 else _extract_numeric_literal(kwargs.get("alpha"))
+    if (beta is not None and float(beta) != 1.0) or (alpha is not None and float(alpha) != 1.0):
+        ctx.fail(f"unsupported scaled addmm (beta/alpha != 1) for {torch_op}: {node.args!r} {kwargs!r}")
     ir_node = IRNode(
         id=node_id(node),
         op="addmm",