rzou 0f768c7866 Barebones flat_apply HOP (#146060) ... This PR: - adds pytree.register_constant for registering a class to be treated as a constant by torch.compile/torch.fx - adds a very barebones flat_apply HOP. This should be sufficient to get mark_traceable working. A lot more work is necessary to get the custom operator case working (when make_fx sees a custom operator with PyTree arg types, it needs to emit a call to the flat_apply HOP). - I expect the flat_apply HOP to change a lot, I want to ship this in the current state to unblock the mark_traceable and custom ops work. Test Plan: - It's kind of difficult to test the barebones flat_apply HOP "works" so I added a really simple test. Pull Request resolved: https://github.com/pytorch/pytorch/pull/146060 Approved by: https://github.com/StrongerXi, https://github.com/yanboliang ghstack dependencies: #146059		2025-02-01 16:17:48 +00:00
..
_internal	Barebones flat_apply HOP (#146060)	2025-02-01 16:17:48 +00:00
__init__.py
_comparison.py	PEP585 update - torch/testing (#145200)	2025-01-20 22:42:42 +00:00
_creation.py	PEP585 update - torch/testing (#145200)	2025-01-20 22:42:42 +00:00
_utils.py