Pathgetter (pmrf.Pathgetter)

class pmrf.Pathgetter(*paths: tuple[Any, ...] | list[Any])

Bases: Module

(experimental) Retrieves values from a PyTree given JAX tree paths. Supports single paths or multiple paths, symmetric to Attrgetter.

__call__(tree: Any) → Any

Call self as a function.

paths: tuple[tuple[Any, ...], ...]