cottax.pytree.crossings module

Reading, writing and completing a pytree at bare key paths.

read_at(tree: Any, keys: tuple[KeyEntry, ...])

The value at keys, stepped to one key at a time.

Return type:

Any

KEY_KINDS: set[type] = {<class 'cottax.core.maths.mints.Mint'>, <class 'jaxlib._jax.pytree.DictKey'>, <class 'jaxlib._jax.pytree.FlattenedIndexKey'>, <class 'jaxlib._jax.pytree.GetAttrKey'>, <class 'jaxlib._jax.pytree.SequenceKey'>}

what a place’s segment may be.

Type:

Every class registered as a key kind (key_kind)

key_kind[K: type](cls: K, step: Callable[[Any, Any], Any] | None = None)

Registers cls as a key kind, so a key of it is a segment as written; step(key, node) is its O(1) step, the one-level flatten without one.

Return type:

K

is_key_entry(key: object)

Whether key is of a registered key kind.

Return type:

bool

write_at(tree: Any, values: Mapping[tuple[KeyEntry, ...], Any])

tree with the value at each key path in values replaced.

One flatten, with each path as a leaf: the treedef is tree’s shape with a slot at each path, so a path may hold a whole subtree and None stays a place. Refuses a path inside another (it would be written twice) and a path the flatten does not reach (it would be written nowhere, in silence).

Return type:

Any

complete_at(tree: Any, names: Iterable[tuple[KeyEntry, ...]])

tree with every one of names it lacks added, holding None. Only a dict grows a key, a missing dict along the way made empty; any other missing key is refused, as writing refuses it. tree itself is not changed.

Return type:

Any

check_antichain(names: Iterable[tuple[KeyEntry, ...]])

Refuse a name that contains another: naming both writes one place twice.

A write precondition only – two nested names read fine, one value read twice. core.maths.places.comparables_in is the same question over `Place`s; this takes raw key tuples.

Return type:

None

resolves(tree: Any, keys: tuple[KeyEntry, ...])

Whether read_at resolves keys in tree.

Return type:

bool

check_pytree_readable(tree: Any, names: Iterable[tuple[KeyEntry, ...]])

Raise unless read_at resolves every one of names, naming the first key that does not.

Return type:

None