[X] I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)
Motivation
There are many use cases in which users only want to register the pytree node type for temporary use. They want to remove the type registration when they exit a context. The object of the corresponding type should become leaf type again.
Required prerequisites
Motivation
There are many use cases in which users only want to register the pytree node type for temporary use. They want to remove the type registration when they exit a context. The object of the corresponding type should become leaf type again.
For example:
https://github.com/pytorch/pytorch/blob/fb91acd33b5ae8c036a1c2373f1dfe6417ebe7a8/torch/onnx/_internal/fx/dynamo_graph_extractor.py#L32-L111
Solution
No response
Alternatives
No response
Additional context
No response