[AtIndexer] make string key points to regex by default, remove BaseKey

CHANGELOG.md

@@ -35,7 +35,8 @@
 
 - `tree_{mask,unmask}` now accepts only callable `cond` argument.
 - Rename `is_frozen` to `is_masked`
-  - frozen could mean non-trainable array, however the masking is not only for arrays but also for other types that will be hidden from jax transformations.
+  - frozen could mean non-trainable array, however the masking is not only for arrays but also for other types that will be hidden across jax transformations.
+- Remove `re.compile(pattern)` to match `re.Pattern` in `where` argument in `AtIndexer`, instead use string `pattern` directly.
 
 ## V0.11.3
 

docs/API/core.rst

@@ -7,7 +7,7 @@
 .. autoclass:: TreeClass 
     :members:
         at
-.. autoclass:: AtIndexer
+.. autoclass:: at
     :members:
         get,
         set,
@@ -16,10 +16,7 @@
         reduce,
         pluck,
 
-.. autoclass:: at
-.. autoclass:: BaseKey
-    :members:
-        __eq__
+.. autoclass:: AtIndexer
 .. autofunction:: autoinit
 .. autofunction:: leafwise
 .. autofunction:: field

sepes/__init__.py

@@ -15,7 +15,7 @@
 from sepes._src.backend import backend_context
 from sepes._src.code_build import autoinit, field, fields
 from sepes._src.tree_base import TreeClass
-from sepes._src.tree_index import AtIndexer, BaseKey, at
+from sepes._src.tree_index import AtIndexer, at
 from sepes._src.tree_mask import (
     is_masked,
     is_nondiff,
@@ -53,7 +53,6 @@
     # indexing utils
     "AtIndexer",
     "at",
-    "BaseKey",
     # tree utils
     "bcmap",
     "leafwise",

sepes/_src/tree_base.py

@@ -270,8 +270,6 @@ def at(self) -> AtIndexer[Self]:
             - ``int`` for positional indexing for sequences.
             - ``...`` to select all leaves.
             - a boolean mask of the same structure as the tree
-            - ``re.Pattern`` to index all keys matching a regex pattern.
-            - an instance of ``BaseKey`` with custom logic to index a pytree.
             - a tuple of the above types to index multiple keys at same level.
 
         Example: