Home
last modified time | relevance | path

Searched refs:call_translations (Results 1 – 5 of 5) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py457 elif eqn.primitive in call_translations:
459 rule = call_translations[eqn.primitive]
914 call_translations: Dict[core.Primitive, Callable] = {} variable
917 call_translations[xla_call_p] = _xla_call_translation_rule
1441 call_translations[pe.remat_call_p] = _remat_translation_rule # type: ignore
1456 call_translations[core.named_call_p] = _named_call_translation_rule
1464 call_translations[core.call_p] = _call_translation_rule
H A Dsharded_jit.py257 xla.call_translations[sharded_call_p] = _sharded_jit_translation_rule
H A Dpxla.py1246 xla.call_translations[xla_pmap_p] = _pmap_translation_rule
/dports/math/py-jax/jax-0.2.9/jax/experimental/
H A Dmaps.py448 xla.call_translations[xmap_p] = _xmap_translation_rule
/dports/math/py-jax/jax-0.2.9/jax/experimental/jax2tf/
H A Djax2tf.py819 for unexpected in xla.call_translations: # Call primitives are inlined