Searched defs:out_pvals (Results 1 – 4 of 4) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | sharded_jit.py | 476 def _pvals_to_results_handler(nrep, npart, partitions, out_pvals): argument
|
| H A D | pxla.py | 1781 out_pvals, devices, backend): argument
|
| /dports/math/py-jax/jax-0.2.9/jax/ |
| H A D | api.py | 1805 def _lift_linearized(jaxpr, primal_avals, consts, io_tree, out_pvals, *py_args): argument
|
| /dports/math/py-jax/jax-0.2.9/jax/_src/lax/ |
| H A D | control_flow.py | 2657 def type_and_const_convert_jaxpr(jaxpr, out_pvals): argument
|