Move HealDA-v2 Triton kernels from experimental to nn.functional
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.3k
- Forks
- 787
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 27
Description
Tracked from review comment in #1758.
The fused Triton kernels currently registered under the "healda::" torch.library namespace (kernels/obs_tokenizer_film.py, kernels/pixel_attention.py) should move to physicsnemo/nn/functional/ once VideoHealDA graduates out of experimental. This would give them a stable namespace (e.g. "physicsnemo::pixel_attn_fwd") and make them reusable by models outside the healda-v2 subpackage.
Scope when addressing:
- Move
custom_opregistrations tophysicsnemo/nn/functional/ - Update the
"healda::"namespace to"physicsnemo::" - Add op schema docstrings and dtype/device constraints
- Update
PixelCrossAttentionandObsTokenizerFiLMto import from the new location
Prerequisite: VideoHealDA stable through one release cycle after #1758 merges.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
First confirm that VideoHealDA has been stable through one release cycle after #1758. Then read kernels/obs_tokenizer_film.py, kernels/pixel_attention.py, and the PixelCrossAttention and ObsTokenizerFiLM call sites; done means the registrations live under physicsnemo/nn/functional/, use the physicsnemo namespace, document schemas and constraints, and imports work from the new location.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100