Home
last modified time | relevance | path

Searched refs:array_result_handler (Results 1 – 3 of 3) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py114 def array_result_handler(device: Optional[Device], aval: core.ShapedArray): function
123 ShapedArray: array_result_handler,
124 ConcreteArray: array_result_handler,
H A Dpxla.py467 def array_result_handler(sharding_spec, indices, aval: ShapedArray): function
469 pxla_result_handlers[ShapedArray] = array_result_handler
470 pxla_result_handlers[ConcreteArray] = array_result_handler
/dports/math/py-jax/jax-0.2.9/jax/_src/lax/
H A Dlax.py1463 return xla.array_result_handler(None, aval)(*xla.device_put(x))