Add State hook system and refactor nn modules and transforms#127

Merged

chaoming0625 merged 52 commits intomainfrom

Jan 1, 2026

Commits on Dec 31, 2025

feat: add parameter module with transform and regularization support
chaoming0625
committed
refactor: reorganize imports in _compatible_import.py for better clarity and structure
chaoming0625
committed
refactor: update import statements in _regularization.py for consistency
chaoming0625
committed
refactor: update ParaM class to inherit from Module for consistency
chaoming0625
committed
feat: add methods for parameter aggregation and named iteration in Module class
chaoming0625
committed
feat: add methods for child and parameter iteration in Module class
chaoming0625
committed
feat: add TraceContextError class and update imports in related modules
chaoming0625
committed
feat: add extend and insert methods to Sequential model for dynamic module management
chaoming0625
committed
test: add init_all_states tests and clean up imports in _module_test.py
chaoming0625
committed
feat: reorganize module imports and update __all__ definition for improved clarity
chaoming0625
committed
feat: implement global hook system for state management with read, write, and restore hooks
chaoming0625
committed
test: enhance unit tests for HookManager with improved structure and clarity
chaoming0625
committed
refactor: streamline hook management by replacing _hooks_manager with hooks property
chaoming0625
committed
refactor: rename _grad_first_order module to _grad_grad for clarity
chaoming0625
committed
refactor: replace vmap with vmap2 in mapping logic and consolidate flatten function
chaoming0625
committed
refactor: rename global state hook functions for consistency and clarity
chaoming0625
committed
refactor: rename _param_module to _par_module for consistency and clarity; update related references
chaoming0625
committed
refactor: replace value() method calls with cache() for consistency in caching logic
chaoming0625
committed
refactor: update copyright information and enhance Module class documentation for clarity
chaoming0625
committed
refactor: update import statement for consistency and add docstring to cache_par_modules method
chaoming0625
committed
refactor: update cache method return type to Data for clarity
chaoming0625
committed
refactor: enhance Vmap and Vmap2 classes with improved documentation and caching methods
chaoming0625
committed
refactor: update tag handling to support multiple tags and improve type consistency
chaoming0625
committed
refactor: enhance delay mechanism with unified ring buffer and interpolation methods
chaoming0625
committed
refactor: add ParamData class for hierarchical state management and enhance __init__ method
chaoming0625
committed
refactor: implement define_param_data method for hierarchical parameter management
chaoming0625
committed
refactor: improve child name handling in module hierarchy for better readability
chaoming0625
committed
refactor: rename ParaM to ParamM and update related documentation and tests
chaoming0625
committed

Commits on Jan 1, 2026