Home
last modified time | relevance | path

Searched refs:shape_with_layout (Results 1 – 25 of 25) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/tests/
Dclient_library_test_base.cc171 absl::Span<GlobalData* const> arguments, const Shape* shape_with_layout) { in ComputeAndCompareLiteral() argument
173 shape_with_layout)); in ComputeAndCompareLiteral()
179 const Shape* shape_with_layout) { in ComputeAndCompareLiteral() argument
181 error, shape_with_layout)); in ComputeAndCompareLiteral()
275 const Shape* shape_with_layout) { in ComputeAndTransfer() argument
293 return ExecuteAndTransfer(computation, arguments, shape_with_layout); in ComputeAndTransfer()
299 const Shape* shape_with_layout) { in ComputeAndCompareLiteralWithStatus() argument
329 if (shape_with_layout != nullptr) { in ComputeAndCompareLiteralWithStatus()
330 layout_shape = *shape_with_layout; in ComputeAndCompareLiteralWithStatus()
337 shape_with_layout = &layout_shape; in ComputeAndCompareLiteralWithStatus()
[all …]
Dclient_library_test_base.h185 const Shape* shape_with_layout = nullptr);
189 const Shape* shape_with_layout = nullptr);
196 const Shape* shape_with_layout = nullptr);
202 const Shape* shape_with_layout = nullptr);
206 const Shape* shape_with_layout = nullptr);
Dlocal_client_execute_test.cc264 Shape shape_with_layout = ShapeUtil::MakeTupleShape( in XLA_TEST_F() local
269 options.set_result_layout(shape_with_layout); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/
Dlayout_assignment.h285 const Shape& shape_with_layout, int64_t priority);
349 Status SetInstructionLayout(const Shape& shape_with_layout,
353 return SetInstructionLayout(shape_with_layout, instruction, mandatory, dfs,
356 Status SetInstructionLayout(const Shape& shape_with_layout,
374 Status SetOperandLayout(const Shape& shape_with_layout,
377 return SetOperandLayout(shape_with_layout, instruction, operand_no,
380 Status SetOperandLayout(const Shape& shape_with_layout,
576 const Shape& shape_with_layout, HloInstruction* instruction);
Dlayout_assignment.cc235 Status LayoutAssignment::SetOperandLayout(const Shape& shape_with_layout, in SetOperandLayout() argument
258 << ShapeUtil::HumanStringWithLayout(shape_with_layout) in SetOperandLayout()
265 shape_with_layout, /*minor_to_major_only=*/true)) { in SetOperandLayout()
289 OperandLayoutConstraint new_constraint(ShapeLayout(shape_with_layout), in SetOperandLayout()
343 LayoutAssignment* assignment, const Shape& shape_with_layout, in SetResultLayout() argument
345 VLOG(3) << " : " << ShapeUtil::HumanStringWithLayout(shape_with_layout) in SetResultLayout()
348 computation_constraint_.ResetResultLayout(ShapeLayout(shape_with_layout), in SetResultLayout()
394 Status LayoutAssignment::SetInstructionLayout(const Shape& shape_with_layout, in SetInstructionLayout() argument
400 << ShapeUtil::HumanStringWithLayout(shape_with_layout) in SetInstructionLayout()
404 if (!ShapeUtil::Compatible(shape_with_layout, instruction->shape())) { in SetInstructionLayout()
[all …]
Dhlo_cost_analysis_test.cc826 Shape shape_with_layout = ShapeUtil::MakeShape(F32, {2, 3, 4, 5}); in TEST_F() local
828 Shape shape_without_layout = shape_with_layout; in TEST_F()
840 shape_with_layout, HloOpcode::kAdd, c1, broadcast)); in TEST_F()
Dservice.cc986 return_shape = Shape(arg->shape_with_layout()); in TransferToClient()
1104 auto literal = Literal::CreateFromShape(Shape(arg->shape_with_layout())); in TransferFromOutfeed()
/external/tensorflow/tensorflow/compiler/xla/python/
Doutfeed_receiver.cc428 Shape shape_with_layout = builder->GetShape(data).ValueOrDie(); in AddOutfeedToBuilder() local
430 &shape_with_layout, [](Shape* subshape, const ShapeIndex&) { in AddOutfeedToBuilder()
436 << "; shape=" << shape_with_layout.ToString(); in AddOutfeedToBuilder()
441 if (!ShapeUtil::Equal(shape_with_layout, found->second)) { in AddOutfeedToBuilder()
445 shape_with_layout.DebugString(), found->second.DebugString(), in AddOutfeedToBuilder()
449 shape_registry_.insert({consumer_id, shape_with_layout}); in AddOutfeedToBuilder()
461 token = OutfeedWithToken(data, token, shape_with_layout, ""); in AddOutfeedToBuilder()
Dops.cc175 absl::Span<const XlaOp> operands, const Shape& shape_with_layout, in BuildOpsSubmodule()
181 builder, call_target_name, operands, shape_with_layout, in BuildOpsSubmodule()
194 absl::Span<const XlaOp> operands, const Shape& shape_with_layout, in BuildOpsSubmodule() argument
202 builder, call_target_name, operands, shape_with_layout, in BuildOpsSubmodule()
Dxla.cc161 Shape shape_with_layout = shape; in PYBIND11_MODULE() local
163 &shape_with_layout, [](Shape* subshape, const ShapeIndex&) { in PYBIND11_MODULE()
168 literal = std::make_shared<Literal>(shape_with_layout); in PYBIND11_MODULE()
Dxla_client_test.py379 shape_with_layout=xla_client.Shape.array_shape(
404 shape_with_layout=xla_client.Shape.array_shape(
2704 shape_with_layout=xla_client.Shape.scalar_shape(np.dtype(np.float32)),
2761 shape_with_layout=xla_client.Shape.scalar_shape(np.dtype(np.uint32)),
/external/tensorflow/tensorflow/compiler/xla/client/
Dclient.cc41 const Shape* shape_with_layout) { in Transfer() argument
44 if (shape_with_layout != nullptr) { in Transfer()
45 *request.mutable_shape_with_layout() = shape_with_layout->ToProto(); in Transfer()
118 const Shape* shape_with_layout, int64_t replica_id, in TransferFromOutfeed() argument
125 if (shape_with_layout != nullptr) { in TransferFromOutfeed()
126 *request.mutable_shape_with_layout() = shape_with_layout->ToProto(); in TransferFromOutfeed()
Dclient.h134 const Shape* shape_with_layout = nullptr);
160 const Shape* shape_with_layout, int64_t replica_id = 0,
Dexecutable_build_options.cc52 const Shape& shape_with_layout) { in set_result_layout() argument
54 result_layout_ = shape_with_layout; in set_result_layout()
Dxla_builder.cc1854 void XlaBuilder::Outfeed(XlaOp operand, const Shape& shape_with_layout, in Outfeed() argument
1862 if (!LayoutUtil::HasLayout(shape_with_layout)) { in Outfeed()
1866 if (!ShapeUtil::Compatible(*operand_shape, shape_with_layout)) { in Outfeed()
1869 ShapeUtil::HumanStringWithLayout(shape_with_layout), in Outfeed()
1872 *instr.mutable_outfeed_shape() = shape_with_layout.ToProto(); in Outfeed()
1907 const Shape& shape_with_layout, in OutfeedWithToken() argument
1911 if (!LayoutUtil::HasLayout(shape_with_layout)) { in OutfeedWithToken()
1915 if (!ShapeUtil::Compatible(*operand_shape, shape_with_layout)) { in OutfeedWithToken()
1918 ShapeUtil::HumanStringWithLayout(shape_with_layout), in OutfeedWithToken()
1921 return OutfeedWithTokenInternal(operand, token, shape_with_layout, in OutfeedWithToken()
[all …]
Dxla_builder.h645 void Outfeed(XlaOp operand, const Shape& shape_with_layout,
648 const Shape& shape_with_layout,
651 XlaOp operand, XlaOp token, const Shape& shape_with_layout,
658 const Shape& shape_with_layout, const std::string& opaque,
672 const XlaComputation* computation, const Shape& shape_with_layout,
686 const XlaComputation& computation, const Shape& shape_with_layout,
758 const std::optional<Shape>& shape_with_layout = std::nullopt,
902 XlaOp SendToHost(XlaOp operand, XlaOp token, const Shape& shape_with_layout,
1255 friend void Outfeed(XlaOp operand, const Shape& shape_with_layout,
1277 absl::Span<const XlaOp> operands, const Shape& shape_with_layout,
[all …]
Dexecutable_build_options.h56 ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout);
/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/
Dops.pyi70 shape_with_layout: Optional[_Layout] = ...,
77 shape_with_layout: Optional[_Layout] = ...) -> XlaOp: ...
186 shape_with_layout: Shape,
196 shape_with_layout: Shape,
265 shape_with_layout: Shape,
/external/tensorflow/tensorflow/compiler/xla/
Dliteral_util.cc404 Shape shape_with_layout = new_literal.shape(); in ReshapeSlice() local
405 *shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); in ReshapeSlice()
412 IndexUtil::LinearIndexToMultidimensionalIndex(shape_with_layout, i); in ReshapeSlice()
Dliteral.cc861 Literal LiteralBase::Relayout(const Shape& shape_with_layout) const { in Relayout()
862 CHECK(ShapeUtil::Compatible(shape_with_layout, shape())) in Relayout()
863 << "Given shape_with_layout " << ShapeUtil::HumanString(shape_with_layout) in Relayout()
866 Literal result = CreateFromShape(shape_with_layout); in Relayout()
Dxla.proto533 ShapeProto shape_with_layout = 2; field
560 ShapeProto shape_with_layout = 1; field
Dliteral.h335 Literal Relayout(const Shape& shape_with_layout) const;
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h242 XlaOp operand, XlaOp token, const Shape& shape_with_layout,
Dmlir_hlo_builder.cc621 XlaOp operand, XlaOp token, const Shape& shape_with_layout, in OutfeedWithTokenInternal() argument
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc1230 xla::Shape shape_with_layout = xla::ShapeUtil::MakeTupleShape(subshapes); in ExportXlaOp() local
1235 value_map[op] = xla::OutfeedWithToken(operand, token, shape_with_layout, in ExportXlaOp()