cottax.pytree.consumers.json module¶

Registers the pytree place’s and node name’s JSON codecs: a key path written and read, a Mint segment as the value it is.

SEGMENTS: dict[str, Callable[[Any], object]] = {'attr': <class 'jaxlib._jax.pytree.GetAttrKey'>, 'dict': <class 'jaxlib._jax.pytree.DictKey'>, 'flat': <class 'jaxlib._jax.pytree.FlattenedIndexKey'>, 'seq': <class 'jaxlib._jax.pytree.SequenceKey'>}¶: Each tag’s key kind, built from the one value it holds (spelled).

write_segment(key: object)[source]¶

write_segment(key: DictKey)

write_segment(key: GetAttrKey)

write_segment(key: SequenceKey)

write_segment(key: FlattenedIndexKey)

spelled(cls: type, tag: str, value: Callable[[Any], Any], build: Callable[[Any], object] | None = None)[source]¶

A key kind’s JSON spelling: a key of cls written {tag: value(key)}, read back by build (cls itself by default). A tag spells one kind.

write_key(key: object, w: Writer)[source]¶

A segment: a purpose as the value it is (its class, by name), any other by its kind’s spelling.

read_key(data: dict[str, Any], r: Reader)[source]¶

segments_of(path: VarPath | NodePath, cls: type)[source]¶

path’s segments, refused unless it is a cls: another class’s with the same keys would be read back as a cls.

class PytreeCodec[source]¶

A variable is a VarPath, a list of segments.

write_place(place: VarPath, w: Writer)[source]¶

place; a value it holds (a Mint) through w.value.

read_place(data: list, r: Reader)[source]¶

The place data spells; a value it holds through r.value.

class PytreeLabelCodec[source]¶

A node name is a NodePath, a list of segments.

write_label(name: NodePath, w: Writer)[source]¶

name; a value it holds (a Mint) through w.value.

read_label(data: list, r: Reader)[source]¶

The node name data spells; a value it holds through r.value.