Standarize program augmentation APIs and their documentation

[chaoming0625]

This pull request introduces a new flag _been_writen to the State class in the brainstate module and makes several updates to the brainstate/augment/_mapping.py file to improve context handling in graph transformations. The most important changes include adding the _been_writen flag, updating methods to handle this new flag, and enhancing the graph transformation functions to support context tags.

State Class Enhancements:

• brainstate/_state.py: Added _been_writen flag to the State class and updated methods like value, update_from_ref, and copy to utilize this flag. [1] [2] [3] [4] [5] [6] [7] [8]

Graph Transformation Enhancements:

• brainstate/augment/_mapping.py: Modified the MapFn class and related functions to include ctxtag for better context handling during graph transformations. [1] [2] [3] [4] [5]

Testing Enhancements:

• brainstate/augment/_mapping_test.py: Added new test cases to validate the functionality of vmap with the new context handling and _been_writen flag.

Context Management Enhancements:

• brainstate/graph/_graph_context.py: Introduced update_context and modified existing context management classes to support ctxtag. [1] [2]