Searched refs:custom_partial_eval_rules (Results 1 – 4 of 4) sorted by relevance
137 if primitive in custom_partial_eval_rules:138 return custom_partial_eval_rules[primitive](self, *tracers, **params)399 custom_partial_eval_rules: Dict[core.Primitive, Callable] = {} variable
728 pe.custom_partial_eval_rules[fun_jvp_p] = fun_jvp_partial_eval
544 pe.custom_partial_eval_rules[while_p] = _while_partial_eval1109 pe.custom_partial_eval_rules[cond_p] = _cond_partial_eval1877 pe.custom_partial_eval_rules[scan_p] = _scan_partial_eval
1009 pe.custom_partial_eval_rules[outside_call_p] = _outside_call_partial_eval_rule