[Feat] Support dataclass in magi_register_custom_op

[themistbeforedawn]

🗂️ PR Category

• ✨ New Feature
• 🚀 Optimization (performance, memory, etc.)
• 💥 Breaking Change
• 🐛 Bug Fix
• 🛠️ Development / Refactoring
• 📚 Documentation
• 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
• 🧪 Testing

📝 Description

What's new

@magi_register_custom_op now accepts frozen-dataclass parameters (recursively nested), so users can group config / flags as a single @dataclass(frozen=True) while torch.library's schema continues to see only primitives:

@dataclasses.dataclass(frozen=True)
class AttnCfg:
    scale: float
    causal: bool = False

@magi_register_custom_op()
def attn(q: torch.Tensor, k: torch.Tensor, cfg: AttnCfg) -> torch.Tensor:
    ...

The same lower-signature pass also handles (transparent to users):

• Literal[...] / string-Enum annotations → auto-downgraded to str
• Unsupported defaults (mutable, dataclass instance, …) are scrubbed from the lowered signature only; user-facing defaults are preserved
• mutates_args accepts either the dataclass-level name (expands to all Tensor leaves) or any lowered leaf name
• backward_fn returns one grad per original parameter (not per leaf); a whole non-differentiable dataclass arg collapses to a single None
• Strict signature validation: variadic args / missing annotations are rejected up-front with actionable errors

Architecture — 4-slot pipeline

Each registration owns up to 4 named objects:

Slot	Object	Created by	Presence
0	fn	user source	Always
1	lowered_fn	this PR	only when the signature needs lowering
2	torch_registered_op	torch.library	Always
3	magi_exposed_op	this PR	only when dataclass flattening is needed

The naming is deliberately dual: torch_registered_op is registered into torch.library's dispatcher; magi_exposed_op is exposed out of Magi to the user.

Architecture — 3 runtime paths

The slot set produced at registration time selects one of three runtime paths:

1. simple  fn → torch_registered_op — zero per-call overhead; returns the OpOverload directly
2. sig-only-rewrite  fn → lowered_fn → torch_registered_op — e.g. Literal downgrade only
3. dataclass-flatten  fn → lowered_fn → torch_registered_op → magi_exposed_op — the wrapper flattens / unflattens on every call; the underlying OpOverload is accessible via op._magi_torch_registered_op

magi_compiler/_magi_register_custom_op.py is laid out 1:1 against this model — 8 numbered sections grouped into registration-time helpers · runtime helpers · main pipeline.

Tests

83 new tests in tests/api_tests/test_register_custom_op.py cover all three runtime paths, autograd bridging through dataclass inputs, nested dataclasses, Optional / Literal / Enum / dtype / device fields, torch.compile integration, and full error-path coverage.

[jiahy0825]

Make these comments self-explanatory by code instead of writing docs.

Or move these comments into a markdown document that clarifies the design of _register_custom_op.

Both acceptable for me haha~

[jiahy0825]

These comments confuse me. Use a human-readable sentence instead of AI-explained comments.

It takes me a long time to understand what these comments say. Provide simplified comments and code examples. I guess you are trying to say the return follows such rules, just paste the examples below:

original_sig: (q: Tensor, cfg: AttnCfg, mode: Literal['a','b']='a') -> Tensor
lowered_sig: (q: Tensor, cfg__scale: float, cfg__causal: bool, mode: str='a') -> Tensor
param_mapping_tree ≈
[
    ("primitive", "q",    "q",    None),
    ("dataclass", "cfg",  AttnCfg, [
        ("primitive", "scale",  "cfg__scale",  None),
        ("primitive", "causal", "cfg__causal", None),
    ]),
    ("primitive", "mode", "mode", None),
]