Searched refs:PyLocalBuffer (Results 1 – 2 of 2) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | xla.py | 54 PyLocalBuffer = Any variable
1032 device_buffer: Union[PyLocalBuffer, "DeviceConstant"], argument
1033 ) -> Union[PyLocalBuffer, "_DeviceArray"]:
1073 device_buffer: PyLocalBuffer): argument
1351 ) -> Callable[[_DeviceArray], PyLocalBuffer]:
1371 force_fun: Callable[[_DeviceArray], PyLocalBuffer]
|
| H A D | pxla.py | 340 args) -> Sequence[Sequence[xb.xla_client._xla.PyLocalBuffer]]:
443 [List[xb.xla_client._xla.PyLocalBuffer]], Any]:
464 PxlaResultHandler = Callable[..., Callable[[List[xb.xla_client._xla.PyLocalBuffer]], Any]]
507 device_buffers: List[xb.xla_client._xla.PyLocalBuffer] = None, argument
1903 …: Sequence[xb.xla_client.Device], replicate: bool=False) -> List[xb.xla_client._xla.PyLocalBuffer]:
|