Skip to content

Expose the register as public register_constant_folder API (fix for Issue #2507)#2967

Open
RohanMali2003 wants to merge 2 commits into
microsoft:mainfrom
RohanMali2003:constant-folding-registry-issue-2507-fix
Open

Expose the register as public register_constant_folder API (fix for Issue #2507)#2967
RohanMali2003 wants to merge 2 commits into
microsoft:mainfrom
RohanMali2003:constant-folding-registry-issue-2507-fix

Conversation

@RohanMali2003

Copy link
Copy Markdown

Summary

Exposes the internal constant-folding registration decorator as a public API: onnxscript.optimizer.register_constant_folder.

Previously, registering a custom constant-folding rule required importing the protected register function from onnxscript.optimizer._constant_folding directly, which linters flag as a private-member import. This re-exports it under a public, more descriptive name, following the same pattern used in #2773 for FOLDED_FROM_KEY.

Closes #2507

Changes

  • onnxscript/optimizer/__init__.py: re-export _constant_folding.register as register_constant_folder, added to __all__.
  • onnxscript/optimizer/_constant_folding.py: added a docstring to PartialEvaluatorRegistry.register documenting the callback signature, parameters, and a usage example.
  • onnxscript/optimizer/_constant_folding_test.py: added test_register_constant_folder, which registers a folder for a mock op under a dedicated test domain (test.custom), verifies fold_constants invokes it, and cleans up the registration in a finally block.

Scope notes

  • Only the register decorator is exposed; other PartialEvaluatorRegistry internals (e.g. lookup_evaluators) remain private.
  • The callback signature references OptimizerContext/OptimizerState, which are currently internal types. I've documented the expected signature informally in the docstring rather than exposing those types — happy to expose them too if preferred.

Testing

  • pytest onnxscript/optimizer/_constant_folding_test.py -v — all passing
  • pytest --doctest-modules onnxscript/optimizer — all passing
  • pytest tests/optimizer/test_models.py -v — all passing (broader sanity check)
  • lintrunner — clean

Sidenote:

I developed and documented this with AI-assisted tooling; all changes were reviewed, and verified by running the full test suite, doctests, and lintrunner locally before submission.

@RohanMali2003

RohanMali2003 commented Jul 16, 2026 via email

Copy link
Copy Markdown
Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

Constant-Folding Registry

1 participant