def rewrite_callback(
prim: core.Primitive, vals: Sequence[core.Tracer],
params: Dict[str,
Any]) -> Union[core.Tracer, Sequence[core.Tracer]]:
if prim in rules:
return rules[prim](*vals, **params)
return prim.bind(*vals, **params)
def find_callback(
prim: core.Primitive, vals: Sequence[core.Tracer],
params: Dict[str,
Any]) -> Union[core.Tracer, Sequence[core.Tracer]]:
vals = prim.bind(*vals, **params)
_contains_query(vals, queries)
return vals
def _process_primitive(self, primitive: Primitive, flat_inputs, kwargs):
if primitive in InitTrace._rules:
return InitTrace._rules[primitive](self)(flat_inputs, kwargs)
if isinstance(primitive, parametrized):
return self.process_parametrized(primitive, *flat_inputs, **kwargs)
return primitive.bind(*flat_inputs, **kwargs)
def default_process_primitive(
self, primitive: jax_core.Primitive, tracers: List['HarvestTracer'],
params: Dict[str, Any]) -> Union['HarvestTracer', List['HarvestTracer']]:
context = trace_util.get_dynamic_context(self)
vals = [t.val for t in tracers]
if primitive is sow_p:
outvals = context.process_sow(*vals, **params)
return jax_util.safe_map(self.pure, outvals)
outvals = primitive.bind(*vals, **params)
if not primitive.multiple_results:
outvals = [outvals]
out_tracers = jax_util.safe_map(self.pure, outvals)
if primitive.multiple_results:
return out_tracers
return out_tracers[0]
