Searched refs:DeviceArrayProtocol (Results 1 – 3 of 3) sorted by relevance
22 def to_dlpack(x: xla.DeviceArrayProtocol, take_ownership: bool = False): argument
1015 DeviceArrayProtocol = Any variable1023 _CppDeviceArray: DeviceArrayProtocol = xc.Buffer1304 def _device_put_device_array(x: Union[DeviceArrayProtocol, _DeviceArray], device: Optional[Device]): argument1310 …vice_array_to_device(x: Union[DeviceArrayProtocol, _DeviceArray], device: Optional[xc.Device]) -> … argument1335 def _force(x: DeviceArrayProtocol) -> DeviceArrayProtocol: argument
397 def _as_slice_indices(arr: xla.DeviceArrayProtocol, idx: Index) -> Tuple[ argument