oryx.core.interpreters.harvest.HarvestTrace

View source on GitHub

An evaluating trace that dispatches to a dynamic context.

oryx.core.interpreters.harvest.HarvestTrace(
    main: MainTrace, sublevel: Sublevel
) -> None

Attributes
level
main
sublevel

Methods

default_process_primitive

View source

default_process_primitive(
    primitive: jax_core.Primitive,
    tracers: List['HarvestTracer'],
    params: oryx.experimental.matching.jax_rewrite.Bindings
) -> Union['HarvestTracer', List['HarvestTracer']]

full_raise

full_raise(
    val
) -> Tracer

lift

View source

lift(
    val: Value
) -> 'HarvestTracer'

post_process_call

View source

post_process_call(
    call_primitive, out_tracers, params
)

post_process_map

View source

post_process_map(
    call_primitive, out_tracers, params
)

process_call

View source

process_call(
    call_primitive: jax_core.Primitive,
    f: Any,
    tracers: List['HarvestTracer'],
    params: oryx.experimental.matching.jax_rewrite.Bindings
)

process_custom_jvp_call

View source

process_custom_jvp_call(
    primitive, fun, jvp, tracers
)

process_custom_transpose

process_custom_transpose(
    prim, call, tracers, **params
)

process_custom_vjp_call

View source

process_custom_vjp_call(
    primitive, fun, fwd, bwd, tracers, out_trees
)

process_map

View source

process_map(
    call_primitive: jax_core.Primitive,
    f: Any,
    tracers: List['HarvestTracer'],
    params: oryx.experimental.matching.jax_rewrite.Bindings
)

process_primitive

View source

process_primitive(
    primitive: jax_core.Primitive,
    tracers: List['HarvestTracer'],
    params: oryx.experimental.matching.jax_rewrite.Bindings
) -> Union['HarvestTracer', List['HarvestTracer']]

process_xmap

process_xmap(
    call_primitive, f, tracers, params
)

pure

View source

pure(
    val: Value
) -> 'HarvestTracer'

sublift

View source

sublift(
    tracer: 'HarvestTracer'
) -> 'HarvestTracer'