Mutates the XlaTensor to set the ShapedBuffer.
| 70 | } |
| 71 | // Mutates the XlaTensor to set the ShapedBuffer. |
| 72 | void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { |
| 73 | shaped_buffer_ = |
| 74 | absl::make_unique<xla::ScopedShapedBuffer>(std::move(shaped_buffer)); |
| 75 | } |
| 76 | |
| 77 | // Some tensors on the device may have known values on the host. We use these |
| 78 | // in on-demand mode to avoid re-copying values from the device if we know the |