cottax.consumers.execution.jax_jit module

JitGraph: a callable graph that can be a jit cache key. Traced: the run traced once by jax, nothing computed: the graph’s jaxpr and every place’s shape.

check_hashable[N: Label, V: Place](name: N, node: NodeDefinition[V])

The node hashes. An array in the graph does not, and is refused: opaque to every query, it is an input the graph does not name. The leaves that fail are named.

Return type:

Iterator[Violation]

check_reassembles[N: Label, V: Place](name: N, node: NodeDefinition[V])

The implementation compares by value, so the same graph built twice is one cache key. A bound method compares its receiver by identity, and a functools.partial (jax.tree_util.Partial is one) defines no __eq__.

Return type:

Iterator[Violation]

class JitGraph[N: Label, V: Place](graph: NestingGraph[N, V])

Bases: CallableGraph[N, V], Generic[N, V]

A callable graph whose every node is a jit cache key.

checks_node: PerLevel[tuple[Callable[[N, NodeDefinition[V]], Iterator[Violation]], ...]] = (<function check_hashable>, <function check_reassembles>)

the one thing a level below adds.

Type:

(name, node) -> violations

check_traced_callable(traced: Traced)

Return type:

Iterator[Violation]

check_traces(traced: Traced)

The run traces: the first refusal it raised, or the node whose body failed.

Return type:

Iterator[Violation]

class Traced[N: Label, V: Place](valued: GraphWithValues[CallableGraph[N, V], V])

Bases: Frozen, Generic[N, V]

The run traced once by jax, nothing computed: the whole graph as one jaxpr, and every place’s shape. Refused: the first refusal the run raises (a body’s miscount, ReturnedUnlike), or a body failing at these shapes (TraceFailed), by node. A value handed in may be an array or a jax.ShapeDtypeStruct; any other value reaches the bodies as it is, so a setting stays a setting.

traced = Traced(GraphWithValues(graph, values)) traced.shapes, traced.jaxpr

valued: GraphWithValues[CallableGraph[N, V], V]

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

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

property jaxpr: Any

The whole graph as one closed jaxpr, its inputs the shapes handed in.

property shapes: OrderedMap[V, Any]

the inputs, then what the graph computes.

Type:

Every place’s shape and dtype (a pytree of them where its value is one)