| /external/tensorflow/tensorflow/compiler/jit/ |
| xla_tensor.cc | 50 xla::Shape on_device_shape = local
54 xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,
59 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first);
|
| xla_launch_util.cc | 238 const xla::Shape on_device_shape = local
240 if (on_device_shape.IsTuple()) {
245 CHECK(xla::ShapeUtil::Equal(shape, on_device_shape))
247 << xla::ShapeUtil::HumanStringWithLayout(on_device_shape)
252 /*on_host_shape=*/shape, /*on_device_shape=*/shape,
270 << output.on_device_shape().DebugString();
281 xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_device_shape()}),
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| shaped_buffer.h | 45 ShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape,
66 const Shape& on_device_shape() const { return on_device_shape_; } function in class:xla::ShapedBuffer
118 // The tree of device buffers. Its shape is on_device_shape().
137 const Shape& on_device_shape,
|
| transfer_manager.cc | 158 const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); local
159 TF_RET_CHECK(on_device_shape.IsArray())
162 << " is not an array: " << ShapeUtil::HumanString(on_device_shape);
163 if (dest.size() < GetByteSizeRequirement(on_device_shape)) {
167 dest.size(), GetByteSizeRequirement(on_device_shape));
169 ShapedBuffer shaped_buffer(/*on_host_shape=*/literal.shape(), on_device_shape,
244 device_buffer.on_device_shape(),
269 TF_RET_CHECK(device_buffer.on_device_shape().IsTuple());
271 TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) ==
276 i < ShapeUtil::TupleElementCount(device_buffer.on_device_shape()); ++i)
[all...] |
| /external/tensorflow/tensorflow/compiler/xrt/ |
| xrt_state.cc | 101 xla::Shape on_device_shape = transfer_manager->HostShapeToDeviceShape(shape); local
104 << xla::ShapeUtil::HumanStringWithLayout(on_device_shape);
110 shape, on_device_shape, allocator, device_ordinal);
113 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first);
170 const xla::Shape& on_device_shape)
174 on_device_shape_(on_device_shape),
203 shaped_buffer.on_device_shape());
216 shaped_buffer.on_device_shape());
262 const xla::Shape& XRTTupleAllocation::on_device_shape() { function in class:tensorflow::XRTTupleAllocation
303 xla::ShapeUtil::TryGetSubshape(parent->on_device_shape(), subshape))
[all...] |
