Skip to content

fix: harden exporter shape and data handling - #222

Merged
JulienBalianSonos merged 5 commits into
mainfrom
fix/exporter-ir-robustness
Jun 30, 2026
Merged

fix: harden exporter shape and data handling#222
JulienBalianSonos merged 5 commits into
mainfrom
fix/exporter-ir-robustness

Conversation

@JulienBalianSonos

Copy link
Copy Markdown
Collaborator

Summary

  • handle HF configs where head_dim exists but is None
  • allow split_with_sizes sections of length zero while still rejecting negative sections
  • make TensorVariable equality safe for tensor subclasses, meta tensors, and numpy arrays
  • add shape-inference rules for aten::_grouped_mm and aten::gather

Quality review

  • This is intentionally separate from the MoE/HQQ export PR. These are generic exporter/IR robustness fixes and would make the MoE PR harder to review if folded in.
  • The head_dim change overlaps behavior already present in PR [RDY] feat: add basic MoE support #55, but keeping it here is reasonable because it is independent and small. If [RDY] feat: add basic MoE support #55 lands first, this will either disappear in rebase or be a no-op conflict.
  • The split change is narrow: only zero-sized sections are newly accepted; negative sizes still fail explicitly.
  • The equality helper now avoids the previous ambiguous/unsafe (data == data).all() path for tensor subclasses and arrays. Base tensors still compare actual values; tensor subclasses compare metadata to avoid forcing subclass-specific equality semantics.
  • The grouped-mm/gather infer rules avoid executing kernels during shape inference, which is safer for optional/private ATen ops and avoids overload-resolution surprises.

Tests

  • uv run ruff format --check packages/llm/torch_to_nnef_llm/config.py packages/llm/tests/test_config.py torch_to_nnef/op/aten/split.py torch_to_nnef/torch_graph/ir_data.py torch_to_nnef/torch_graph/ir_op.py tests/test_grouped_mm_infer.py tests/test_ir_data_equality.py tests/test_aten_parser_fixes.py
  • uv run ruff check packages/llm/torch_to_nnef_llm/config.py packages/llm/tests/test_config.py torch_to_nnef/op/aten/split.py torch_to_nnef/torch_graph/ir_data.py torch_to_nnef/torch_graph/ir_op.py tests/test_grouped_mm_infer.py tests/test_ir_data_equality.py tests/test_aten_parser_fixes.py
  • uv run pytest packages/llm/tests/test_config.py -q
  • uv run pytest tests/test_grouped_mm_infer.py tests/test_ir_data_equality.py tests/test_aten_parser_fixes.py -q

@JulienBalianSonos
JulienBalianSonos force-pushed the fix/exporter-ir-robustness branch from 2a4c116 to 45e6257 Compare June 30, 2026 10:14
@JulienBalianSonos
JulienBalianSonos merged commit f22e00d into main Jun 30, 2026
39 checks passed
@JulienBalianSonos
JulienBalianSonos deleted the fix/exporter-ir-robustness branch June 30, 2026 11:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant