Quark Create Shapeshifter Pass

quark-create-shapeshifter-pass

Purpose

Author a new ShapeShifter transformation pass that is correct on the first try. ShapeShifter is Quark's pass-based graph-transformation framework (quark/shapeshifter/): each pass is a self-contained unit that takes a model, applies one transformation, and returns it. Passes auto-register by filename via a @register_pass decorator, so a new pass is usable from the CLI, the Python API, and the ONNX quantizer's ShapeShifterYaml with zero manual wiring — if it follows the naming, subclassing, and config conventions exactly. This skill replays those conventions in authoring order and flags the non-obvious traps (defaults are not auto-applied; class/file naming is load-bearing) so the developer focuses on the transformation logic.

Inputs

• Required: the transformation the pass performs (one sentence); the backend it targets (ONNX onnx.ModelProto or PyTorch callable — a pass is one or the other, never both).
• Required: a pass name in snake_case, prefixed onnx_ or pytorch_ (this becomes the filename, the registry key, and the YAML key — pick it carefully).
• Optional: config parameters the pass exposes (name, type, default, whether required); whether it is a core pass (quark/shapeshifter/passes/) or a community pass (quark/contrib/shapeshifter_community_passes/); a minimal model that exercises it (for the test).

Outputs: shapeshifter_pass.py

The primary artifact is the new pass module, written to quark/shapeshifter/passes/<pass_name>.py (core) or quark/contrib/shapeshifter_community_passes/<pass_name>.py (community). <pass_name> matches the chosen name (e.g. onnx_drop_identity.py). Side-effect artifacts:

• test/test_for_cli/test_shapeshifter_<pass_name>.py — one test per pass (project convention).
• A documentation entry in docs/source/quark_shapeshifter_onnx_passes.rst (ONNX) or docs/source/quark_shapeshifter_torch_passes.rst (PyTorch).

No JSON schema — this artifact is Quark source code, not a cross-skill handoff artifact.

Interaction Flow

1. Intake — capture the transformation, the backend, the pass name, and any config params. If the developer described the transform in prior conversation, extract these first.
2. Route — confirm backend (ONNX vs PyTorch) and core-vs-community placement; derive the file path, class name, and YAML key from the name.
3. Plan — present the file path, class skeleton, _default_config params, and the _run_for_config outline before writing. Confirm the name does not collide with an existing registry entry (ls quark/shapeshifter/passes/).
4. Confirm — get explicit approval before writing source files under quark/.
5. Execute or Summarize — write the pass, the test, and the doc entry; run the test and confirm the pass registers. Summarize what was created and how to invoke it.

Backend & Placement Decision

A workflow must be all-ONNX or all-PyTorch; the Engine raises ValueError on mixing.

Naming Rules (load-bearing)

• Filename = pass name = registry key = YAML key. onnx_drop_identity.py → pass name onnx_drop_identity, used verbatim under passes: in YAML. There is no separate name string.
• Prefix onnx_ or pytorch_.
• Files starting with _ (e.g. __init__.py) are skipped by discovery.
• Class name must end in Pass (only such classes are discovered / exported). Convention: ONNX<Thing>Pass / Pytorch<Thing>Pass.
• Duplicate pass names raise ValueError at import (with a core-vs-community conflict message).

Authoring the Pass

Two required methods (see quark/shapeshifter/pass_base.py):

• _default_config(self) -> dict[str, PassConfigParam] — declare each config param with PassConfigParam(type_, default_value, required, description); end with config.update(self.config) (the convention every pass follows).
• _run_for_config(self, model, config) -> model — the transformation. ONNX: takes/returns onnx.ModelProto. PyTorch: takes/returns any Callable.

ONNX skeleton (quark/shapeshifter/passes/onnx_drop_identity.py):

#
# Copyright (C) 2025 - 2026 Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: MIT
#

from typing import Any

import onnx
from onnx import ModelProto
from onnxruntime.quantization.onnx_model import ONNXModel  # helpers: remove_nodes, etc.

from quark.common.utils.log import ScreenLogger
from quark.shapeshifter.pass_base import ONNXPass, register_pass
from quark.shapeshifter.pass_config import PassConfigParam

logger = ScreenLogger(__name__)


@register_pass
class ONNXDropIdentityPass(ONNXPass):
    """Remove Identity nodes from the graph."""

    def _default_config(self) -> dict[str, PassConfigParam]:
        config = {
            "drop_identity": PassConfigParam(
                type_=bool,
                default_value=True,
                required=True,
                description="Whether to remove Identity nodes.",
            ),
        }
        config.update(self.config)
        return config

    def _run_for_config(self, model: ModelProto, config: dict[str, Any]) -> ModelProto:
        # config is the RAW user dict from YAML — read values directly and
        # supply your own default (see Critical Gotcha below).
        if not config.get("drop_identity", True):
            logger.warning("onnx_drop_identity: drop_identity is False, skipping.")
            return model
        onnx_model = ONNXModel(model)
        # ... transform onnx_model.model, then clean up ...
        onnx_model.topological_sort()
        return onnx_model.model

PyTorch skeleton mirrors this: subclass PytorchPass, _run_for_config(self, model, config) returns the (possibly mutated) callable. See quark/shapeshifter/passes/pytorch_remove_dropout.py.