Searched refs:AsShapedBuffer (Results 1 – 5 of 5) sorted by relevance
51 TEST(TrackedDeviceBufferTest, AsShapedBuffer) { in TEST() argument67 ShapedBuffer shaped_a = a_buffer->AsShapedBuffer( in TEST()69 ShapedBuffer shaped_b = b_buffer->AsShapedBuffer( in TEST()71 ShapedBuffer shaped_c = c_buffer->AsShapedBuffer( in TEST()
141 ShapedBuffer AsShapedBuffer(const Shape& on_device_shape) const;
778 ShapedBuffer buffer = device_buffer->AsShapedBuffer(on_device_shape); in BufferFromHostBuffer()901 ShapedBuffer buffer = device_buffer->AsShapedBuffer(on_device_shape); in BufferFromHostLiteral()1236 ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(on_device_shape_); in ToLiteral()1267 StatusOr<ShapedBuffer> PjRtStreamExecutorBuffer::AsShapedBuffer() const { in AsShapedBuffer() function in xla::PjRtStreamExecutorBuffer1273 return device_buffer_->AsShapedBuffer(on_device_shape_); in AsShapedBuffer()1303 TF_ASSIGN_OR_RETURN(ShapedBuffer src_buffer, AsShapedBuffer()); in CopyToDeviceHelper()1309 ShapedBuffer dst_buffer = dst_device_buffer->AsShapedBuffer(on_device_shape_); in CopyToDeviceHelper()
120 ShapedBuffer TrackedDeviceBuffer::AsShapedBuffer( in AsShapedBuffer() function in xla::TrackedDeviceBuffer
506 StatusOr<ShapedBuffer> AsShapedBuffer() const;