Searched refs:on_device_shape_ (Results 1 – 6 of 6) sorted by relevance
69 const Shape& on_device_shape() const { return on_device_shape_; } in on_device_shape()93 CHECK(ShapeUtil::Equal(buffers.shape(), on_device_shape_)); in set_buffers()95 buffers_.replace_shape_ptr(&on_device_shape_); in set_buffers()102 CHECK(ShapeUtil::EqualStructure(on_device_shape, on_device_shape_)) in set_shapes()104 << ", old: " << on_device_shape_; in set_shapes()106 on_device_shape_ = on_device_shape; in set_shapes()107 buffers_.replace_shape_ptr(&on_device_shape_); in set_shapes()130 Shape on_device_shape_; variable
35 : on_device_shape_(std::move(on_device_shape)), in ShapedBuffer()37 buffers_(&on_device_shape_) { in ShapedBuffer()38 on_host_shape_ = ShapeUtil::DeviceShapeToHostShape(on_device_shape_); in ShapedBuffer()47 on_device_shape_(std::move(s.on_device_shape_)), in ShapedBuffer()53 buffers_.replace_shape_ptr(&on_device_shape_); in ShapedBuffer()57 on_device_shape_ = std::move(s.on_device_shape_); in operator =()64 buffers_.replace_shape_ptr(&on_device_shape_); in operator =()
469 const Shape& on_device_shape() const override { return on_device_shape_; } in on_device_shape()475 return on_device_shape_.IsTuple() && in IsEmptyTuple()476 on_device_shape_.tuple_shapes_size() == 0; in IsEmptyTuple()588 const Shape on_device_shape_; variable
630 if (on_device_shape_.IsTuple()) { in ReleaseDeviceMemoryOwnership()1001 on_device_shape_(std::move(on_device_shape)), in PjRtStreamExecutorBuffer()1022 ->GetByteSizeRequirement(on_device_shape_); in OnDeviceSizeInBytes()1236 ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(on_device_shape_); in ToLiteral()1273 return device_buffer_->AsShapedBuffer(on_device_shape_); in AsShapedBuffer()1299 ShapeUtil::DeviceShapeToHostShape(on_device_shape_), in CopyToDeviceHelper()1309 ShapedBuffer dst_buffer = dst_device_buffer->AsShapedBuffer(on_device_shape_); in CopyToDeviceHelper()
154 on_device_shape_(on_device_shape), in XRTTupleAllocation()155 buffers_(&on_device_shape_), in XRTTupleAllocation()353 return on_device_shape_; in on_device_shape()
284 const xla::Shape on_device_shape_; variable