diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py index 4cb8967f8..90cb587ec 100644 --- a/python/cactus/transpile/aten_ops.py +++ b/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 diff --git a/python/cactus/transpile/capture_jax.py b/python/cactus/transpile/capture_jax.py index 7033f74ca..bef4b50e3 100644 --- a/python/cactus/transpile/capture_jax.py +++ b/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) diff --git a/python/cactus/transpile/fusion/rms_norm.py b/python/cactus/transpile/fusion/rms_norm.py index 8951e3c42..322833744 100644 --- a/python/cactus/transpile/fusion/rms_norm.py +++ b/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 diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index 17ed8d2a7..d606d32f5 100644 --- a/python/cactus/transpile/importers.py +++ b/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",