Searched refs:TpuStream (Results 1 – 4 of 4) sorted by relevance
29 class TpuStream : public tensorflow::tpu::TpuStreamInterface {33 explicit TpuStream(SE_Stream* stream) : stream_(stream) {} in TpuStream() function34 ~TpuStream() override { in ~TpuStream()42 stream_, static_cast<TpuStream*>(other)->stream_); in IsSameSharedMemoryLocation()
179 auto ptr = absl::make_unique<tpu::TpuStream>(tpu_stream); in GetStreamImplementation()
47 static_cast<tensorflow::tpu::TpuStream*>( in ToC()77 static_cast<tensorflow::tpu::TpuStream*>(impl)->se_stream(); in ToC()
69 auto* transfer_tpu_stream = tensorflow::down_cast<tf_tpu::TpuStream*>( in ThenMemcpyDeviceToDevice()