GitHub

oryx.core.interpreters.propagate.propagate

Propagates cells in a Jaxpr using a set of rules.

oryx.core.interpreters.propagate.propagate(
    cell_type: Type[oryx.core.interpreters.propagate.Cell],
    rules: Dict[jax_core.Primitive, PropagationRule],
    jaxpr: pe.Jaxpr,
    constcells: List[oryx.core.interpreters.propagate.Cell],
    incells: List[oryx.core.interpreters.propagate.Cell],
    outcells: List[oryx.core.interpreters.propagate.Cell]
) -> oryx.core.interpreters.propagate.Environment

Args
`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

Returns
The Jaxpr environment after propagation has terminated

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.

Last updated 2021-01-29 UTC.

oryx.core.interpreters.propagate.propagate

Args

Returns