quark-create-shapeshifter-pass

Files (1)

• SKILL.md

SKILL.md (first lines)

# 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

| Question | Choose |
|----------|--------|
| Operates on `onnx.ModelProto` (a graph)? | ONNX pass → subclass `ONNXPass`, prefix `onnx_` |
| Operates on a PyTorch callable (`nn.Module`, function)? | PyTorch pass → subclass `PytorchPass`, prefix `pytorch_` |
| Officially maintained / production? | Core → `quark/shapeshifter/passes/` |
| Contributed / experimental? | Community → `quark/contrib/shapeshifter_community_passes/` |

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`):

```python
#
# 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

From amd/Quark, MIT.

More from amd/Quark