Extend capabilities of tree_get, tree_set.

1. Enable tree_get, tree_set to filter for the name of a named tuple in the path to a key (hence filter the name of a state in a chained transformation).
This enables distinguishing for attributes identical in two different states except that the names of the states are different.

2. Enable tree_get, tree_set to fetch or set named tuples by the name of the name tuple. This is handy to fetch a given state in the overall state of a chained optimizer.

PiperOrigin-RevId: 617976030

docs/api/utilities.rst

@@ -95,6 +95,7 @@ Tree
 .. currentmodule:: optax.tree_utils
 
 .. autosummary::
+    NamedTupleKey
     tree_add
     tree_add_scalar_mul
     tree_div
@@ -112,6 +113,10 @@ Tree
     tree_vdot
     tree_zeros_like
 
+NamedTupleKey
+~~~~~~~~~~~~~
+.. autoclass:: tree_add
+
 Tree add
 ~~~~~~~~
 .. autofunction:: tree_add

optax/tree_utils/__init__.py

@@ -14,11 +14,13 @@
 # ==============================================================================
 """The tree_utils sub-package."""
 
+from optax.tree_utils._state_utils import NamedTupleKey
 from optax.tree_utils._state_utils import tree_get
 from optax.tree_utils._state_utils import tree_get_all_with_path
 from optax.tree_utils._state_utils import tree_map_params
 from optax.tree_utils._state_utils import tree_set
 
+
 from optax.tree_utils._tree_math import tree_add
 from optax.tree_utils._tree_math import tree_add_scalar_mul
 from optax.tree_utils._tree_math import tree_div