Searched refs:_xla_callable (Results 1 – 2 of 2) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | xla.py | 261 return _xla_callable(lu.wrap_init(prim_fun), device, None, "prim", donated_invars,
580 compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
649 def _xla_callable(fun: lu.WrappedFun, device, backend, name, donated_invars, *arg_specs): function
|
| /dports/math/py-jax/jax-0.2.9/jax/ |
| H A D | api.py | 305 xla._xla_callable.most_recent_entry())
|