Home
last modified time | relevance | path

Searched refs:shape_tree (Results 1 – 11 of 11) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/
Dshape_tree_test.cc121 ShapeTree<int> shape_tree{ShapeUtil::MakeTupleShape({})}; in TEST_F() local
122 EXPECT_EQ(0, shape_tree.leaf_count()); in TEST_F()
126 ShapeTree<int> shape_tree{array_shape_}; in TEST_F() local
127 *shape_tree.mutable_element({}) = 42; in TEST_F()
128 EXPECT_EQ(42, shape_tree.element({})); in TEST_F()
129 *shape_tree.mutable_element({}) = 123; in TEST_F()
130 EXPECT_EQ(123, shape_tree.element({})); in TEST_F()
132 EXPECT_TRUE(ShapeUtil::Compatible(array_shape_, shape_tree.shape())); in TEST_F()
135 ShapeTree<int> copy{shape_tree}; in TEST_F()
141 EXPECT_EQ(123, shape_tree.element({})); in TEST_F()
[all …]
DBUILD684 name = "shape_tree",
685 hdrs = ["shape_tree.h"],
703 ":shape_tree",
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_test.cc72 ShapeTree<HloSharding> shape_tree = in TEST_F() local
74 EXPECT_EQ(shape_tree.element({}), sharding); in TEST_F()
75 EXPECT_TRUE(shape_tree.IsLeaf({})); in TEST_F()
142 ShapeTree<HloSharding> shape_tree = in TEST_F() local
144 EXPECT_EQ(shape_tree.element({0}), HloSharding::Replicate()); in TEST_F()
145 EXPECT_EQ(shape_tree.element({1, 0}), HloSharding::AssignDevice(0)); in TEST_F()
146 EXPECT_EQ(shape_tree.element({2}), tiled_sharding); in TEST_F()
191 ShapeTree<HloSharding> shape_tree(ShapeUtil::MakeTupleShape({}), in TEST_F() local
194 HloSharding sharding2 = HloSharding::Tuple(shape_tree); in TEST_F()
199 ShapeTree<HloSharding> shape_tree(ShapeUtil::MakeTupleShape({}), in TEST_F() local
[all …]
Dhlo_sharding.cc280 ShapeTree<HloSharding> shape_tree = GetAsShapeTree(shape); in ValidateTuple() local
281 for (const auto& index_to_sharding : shape_tree.leaves()) { in ValidateTuple()
DBUILD143 "//tensorflow/compiler/xla:shape_tree",
358 "//tensorflow/compiler/xla:shape_tree",
858 "//tensorflow/compiler/xla:shape_tree",
914 "//tensorflow/compiler/xla:shape_tree",
2345 "//tensorflow/compiler/xla:shape_tree",
2415 "//tensorflow/compiler/xla:shape_tree",
2455 "//tensorflow/compiler/xla:shape_tree",
2537 "//tensorflow/compiler/xla:shape_tree",
3844 "//tensorflow/compiler/xla:shape_tree",
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_transfer_manager.cc121 ShapeTree<std::unique_ptr<gpu::OutfeedBuffer>>* shape_tree) { in ShapeTreeToLiteral() argument
125 ShapeTree<std::unique_ptr<gpu::OutfeedBuffer>>* shape_tree, in ShapeTreeToLiteral()
127 const Shape& shape = ShapeUtil::GetSubshape(shape_tree->shape(), *index); in ShapeTreeToLiteral()
129 (*shape_tree->mutable_element(*index))->WaitUntilAvailable(); in ShapeTreeToLiteral()
138 helper(shape_tree, index); in ShapeTreeToLiteral()
144 Helper::helper(shape_tree, &index); in ShapeTreeToLiteral()
Dhlo_to_ir_bindings.cc281 const ShapeTree<llvm::Value*>& shape_tree = it->second; in ToString() local
283 const llvm::Value* val = shape_tree.begin()->second; in ToString()
289 for (auto shape_it = shape_tree.begin(); shape_it != shape_tree.end(); in ToString()
DBUILD382 "//tensorflow/compiler/xla:shape_tree",
787 "//tensorflow/compiler/xla:shape_tree",
915 "//tensorflow/compiler/xla:shape_tree",
929 "//tensorflow/compiler/xla:shape_tree",
/external/tensorflow/tensorflow/compiler/xla/client/
DBUILD184 "//tensorflow/compiler/xla:shape_tree",
/external/tensorflow/tensorflow/compiler/xla/python/
Dlocal_computation_builder.cc182 ShapeTree<se::DeviceMemoryBase>& shape_tree = tuple_buffer.buffers(); in DestructureTuple() local
195 shape_tree.mutable_element(original); in DestructureTuple()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
DBUILD231 "//tensorflow/compiler/xla:shape_tree",