Searched refs:_pval_to_result_handler (Results 1 – 3 of 3) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | sharded_jit.py | 474 global _pvals_to_results_handler, _pval_to_result_handler
478 handlers = [_pval_to_result_handler(npart, parts, out_pval)
493 def _pval_to_result_handler(npart, parts, pval): function
|
| H A D | pxla.py | 762 _pval_to_result_handler( # type: ignore
1771 _pvals_to_results_handler, _pval_to_result_handler, replicate, \
1786 _pval_to_result_handler(size, nrep, npart, parts, pval, devices, backend)
1801 def _pval_to_result_handler(axis_size, nrep, npart, parts, pval, devices, backend): function
|
| H A D | xla.py | 673 result_handlers = map(partial(_pval_to_result_handler, device), pvals) # type: ignore
1469 global _pval_to_result_handler
1471 def _pval_to_result_handler(device, pval): function
|