oryx.core.interpreters.propagate.propagate

Propagates cells in a Jaxpr using a set of rules.

cell_type used to instantiate literals into cells
rules maps JAX primitives to propagation rule functions
jaxpr used to construct the propagation graph
constcells used to populate the Jaxpr's constvars
incells used to populate the Jaxpr's invars
outcells used to populate the Jaxpr's outcells

The Jaxpr environment after propagation has terminated