Searched refs:array_result_handler (Results 1 – 3 of 3) sorted by relevance
114 def array_result_handler(device: Optional[Device], aval: core.ShapedArray): function123 ShapedArray: array_result_handler,124 ConcreteArray: array_result_handler,
467 def array_result_handler(sharding_spec, indices, aval: ShapedArray): function469 pxla_result_handlers[ShapedArray] = array_result_handler470 pxla_result_handlers[ConcreteArray] = array_result_handler
1463 return xla.array_result_handler(None, aval)(*xla.device_put(x))