Python _xla_call_translation_rule Examples

Example #1

def _nest_translation_rule(*args, backend, name, call_jaxpr, scope, **_):
    return xla._xla_call_translation_rule(  # pylint: disable=protected-access
        *args,
        name=jax_util.wrap_name(name, f'nest[{scope}]'),
        backend=backend,
        call_jaxpr=call_jaxpr,
        donated_invars=(False, ) * len(args))

Example #2

 def _translation_rule(*args, backend, name, call_jaxpr, **kwargs):
     del kwargs
     return xla._xla_call_translation_rule(
         *args,
         name=name,
         backend=backend,  # pylint: disable=protected-access
         call_jaxpr=call_jaxpr,
         donated_invars=(False, ) * len(args))

Example #3

File: primitive.py Project: rupei/probability

 def rule(*args, backend, name, call_jaxpr, **params):
   new_params = dict(name=name, backend=backend, call_jaxpr=call_jaxpr)
   new_params['donated_invars'] = params.get('donated_invars',
                                             (False,) * len(args))
   return xla._xla_call_translation_rule(*args, **new_params)  # pylint: disable=protected-access