Searched refs:call_translations (Results 1 – 5 of 5) sorted by relevance
457 elif eqn.primitive in call_translations:459 rule = call_translations[eqn.primitive]914 call_translations: Dict[core.Primitive, Callable] = {} variable917 call_translations[xla_call_p] = _xla_call_translation_rule1441 call_translations[pe.remat_call_p] = _remat_translation_rule # type: ignore1456 call_translations[core.named_call_p] = _named_call_translation_rule1464 call_translations[core.call_p] = _call_translation_rule
257 xla.call_translations[sharded_call_p] = _sharded_jit_translation_rule
1246 xla.call_translations[xla_pmap_p] = _pmap_translation_rule
448 xla.call_translations[xmap_p] = _xmap_translation_rule
819 for unexpected in xla.call_translations: # Call primitives are inlined