cottax.pytree.path module

PyTree path: KeyPath from Jax with KeyEntries.

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

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, ...])

Bases: Path

Where a value lives.

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

Bases: Path

Where a computation lives.