oryx.core.interpreters.harvest.HarvestTrace

An evaluating trace that dispatches to a dynamic context.

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

level

main

sublevel

Methods

default_process_primitive

View source

default_process_primitive(
    primitive: jax_core.Primitive,
    tracers: List[HarvestTracer],
    params: Dict[str, Any]
) -> 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_custom_jvp_call

View source

post_process_custom_jvp_call(
    out_tracers, jvp_was_run
)

post_process_custom_vjp_call

View source

post_process_custom_vjp_call(
    out_tracers, params
)

post_process_custom_vjp_call_fwd

View source

post_process_custom_vjp_call_fwd(
    out_tracers, out_trees
)

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: Dict[str, Any]
)

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: Dict[str, Any]
)

process_primitive

View source

process_primitive(
    primitive: jax_core.Primitive,
    tracers: List[HarvestTracer],
    params: Dict[str, Any]
) -> 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