cottax.pytree.consumers.surface.names module

The pytree place’s names: a node name and a place from each other’s keys, and a key one deeper (PytreeDerivedNames), what a path written in full means (WrittenKeys), and what a bare string means, a jax key under a root (Under) or nothing, every place written in full (InFull).

class PytreeDerivedNames

Bases: DerivedNames[NodePath, VarPath]

A node name and a place give each other their keys; a name one deeper is a string key; a model tree is read by key path.

beside(place: VarPath, name: str)

place one name deeper: where a read written as a value sits, under its purpose.

Return type:

VarPath

node_at(place: VarPath)

The node name a place gives: where a closure over it binds.

Return type:

NodePath

place_at(node: NodePath)

The place a node name gives: where a driver’s per-problem data sits.

Return type:

VarPath

node_in_tree(steps: tuple[Step, ...])

The node name a position in a tree of models gives: the steps to it from the root.

Return type:

NodePath

class WrittenKeys

Bases: Written[NodePath, VarPath]

A path written in full: a V path the place it builds, an N path the node name.

places(written: Keys)

The places written stands for; None where it is no place.

Return type:

tuple[VarPath, ...] | None

node(written: Keys, names: DerivedNames)

The node name written takes, names the surface’s rules; None where it is no node name.

Return type:

NodePath | None

class PytreeLocator

Bases: Locator[VarPath, NodePath]

What every pytree locator decides alike: a string where a node is due is a key at the root, as a slot in a mapping names one. What a string means where a place is due is each final’s.

places

alias of VarPath

nodes

alias of NodePath

node(s: str)

A string written where a node is due.

Return type:

NodePath

key_of(step: Step)

The jax key a step down a model tree is: a field a GetAttrKey, a key a DictKey, an index a SequenceKey.

Return type:

KeyEntry

check_name_key(locator: Under)

Return type:

Iterator[Violation]

class Under(root: Var, key: type[GetAttrKey] | type[DictKey])

Bases: PytreeLocator

A place string one key under a root, the key a GetAttrKey or a DictKey.

root: Var

key: type[GetAttrKey] | type[DictKey]

checks: PerLevel[tuple[Callable[[Any], object], ...]] = (<function check_name_key>,)

Callable [self] -> iterable of violations, which are raised at construction unless skipped (see unchecked).

place(s: str)

A string written where a place is due.

Return type:

VarPath

property V: Keys

any path relative to it, in V’s syntax, where a string can say only one name: tfcoil.locator.V.coils[0].j.

Type:

A builder at the root

class InFull

Bases: PytreeLocator

No place strings: every place written in full, a V path or a VarPath.

place(s: str)

A string written where a place is due.

Return type:

VarPath

property V: Keys

the root is the root.

Type:

V itself

class NotANameKey(key: str)

Bases: Violation

key: str

template: AbstractClassVar[str] = 'a scope writes a name as a `GetAttrKey` or a `DictKey`, not {key}: any other key is written in a `V` path'

e.g. ‘output {place!r} is owned by both {other!r} and {node!r}’

Type:

Wording