cottax.pytree

A place as a jax key path, the one implementation of core.Place cottax ships: VarPath / NodePath (path.py), its derived names, a Mint as the root segment (mint.py), and indexing a pytree by key path (crossings.py), where a key kind of one’s own is registered.

class NodePath(segments: tuple[KeyEntry, ...])[source]

Bases: Path

Where a computation lives.

class Path(segments: tuple[KeyEntry, ...])[source]

Bases: ABC

Place in a Pytree

segments: tuple[KeyEntry, ...]
property spelling: str

.mesh.nodes, [‘eq’][0].psi. What a renderer draws and what str(path) gives.

A label, not an identity: DictKey(0) and SequenceKey(0) both write [0], so nothing keys, orders or compares by it.

Type:

jax’s own spelling, root dot and all

class VarPath(segments: tuple[KeyEntry, ...])[source]

Bases: Path

Where a value lives.

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

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

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

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

Return type:

None

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

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

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

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

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

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

Return type:

Any

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

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

Subpackages

Submodules