| /external/tensorflow/tensorflow/compiler/xla/ |
| D | shape_tree_test.cc | 121 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 …]
|
| D | BUILD | 773 name = "shape_tree",
774 hdrs = ["shape_tree.h"],
792 ":shape_tree",
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_replication_analysis.cc | 287 ShapeTree<bool> shape_tree(inst->shape(), true); in ComputeHloReplicationOnComputation() local
289 shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(i)], {}, {i}); in ComputeHloReplicationOnComputation()
291 changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); in ComputeHloReplicationOnComputation()
293 ShapeTree<bool> shape_tree(inst->shape(), true); in ComputeHloReplicationOnComputation() local
294 shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(0)], in ComputeHloReplicationOnComputation()
296 changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); in ComputeHloReplicationOnComputation()
298 ShapeTree<bool> shape_tree(inst->shape(), false); in ComputeHloReplicationOnComputation() local
301 shape_tree.ForEachMutableElement( in ComputeHloReplicationOnComputation()
306 changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); in ComputeHloReplicationOnComputation()
312 ShapeTree<bool> shape_tree(inst->shape(), true); in ComputeHloReplicationOnComputation() local
[all …]
|
| D | hlo_sharding_test.cc | 87 ShapeTree<HloSharding> shape_tree = in TEST_F() local
89 EXPECT_EQ(shape_tree.element({}), sharding); in TEST_F()
90 EXPECT_TRUE(shape_tree.IsLeaf({})); in TEST_F()
177 ShapeTree<HloSharding> shape_tree = in TEST_F() local
179 EXPECT_EQ(shape_tree.element({0}), HloSharding::Replicate()); in TEST_F()
180 EXPECT_EQ(shape_tree.element({1, 0}), HloSharding::AssignDevice(0)); in TEST_F()
181 EXPECT_EQ(shape_tree.element({2}), tiled_sharding); in TEST_F()
226 ShapeTree<HloSharding> shape_tree(ShapeUtil::MakeTupleShape({}), in TEST_F() local
229 HloSharding sharding2 = HloSharding::Tuple(shape_tree); in TEST_F()
234 ShapeTree<HloSharding> shape_tree(ShapeUtil::MakeTupleShape({}), in TEST_F() local
[all …]
|
| D | hlo_sharding.cc | 388 ShapeTree<HloSharding> shape_tree = GetAsShapeTree(shape); in ValidateTuple() local
389 for (const auto& index_to_sharding : shape_tree.leaves()) { in ValidateTuple()
|
| D | BUILD | 167 "//tensorflow/compiler/xla:shape_tree",
452 "//tensorflow/compiler/xla:shape_tree",
1085 "//tensorflow/compiler/xla:shape_tree",
1145 "//tensorflow/compiler/xla:shape_tree",
2787 "//tensorflow/compiler/xla:shape_tree",
3184 "//tensorflow/compiler/xla:shape_tree",
3342 "//tensorflow/compiler/xla:shape_tree",
3381 "//tensorflow/compiler/xla:shape_tree",
3465 "//tensorflow/compiler/xla:shape_tree",
5056 "//tensorflow/compiler/xla:shape_tree",
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | hlo_to_ir_bindings.cc | 259 const ShapeTree<llvm::Value*>& shape_tree = it->second; in ToString() local
261 const llvm::Value* val = shape_tree.begin()->second; in ToString()
267 for (auto shape_it = shape_tree.begin(); shape_it != shape_tree.end(); in ToString()
|
| D | BUILD | 622 "//tensorflow/compiler/xla:shape_tree",
1213 "//tensorflow/compiler/xla:shape_tree",
1479 "//tensorflow/compiler/xla:shape_tree",
1494 "//tensorflow/compiler/xla:shape_tree",
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_on_demand_compiler.cc | 109 ApiConverter::ToC(arg.shape(), &se_args[i]->shape_tree.shape); in ExecuteAsyncOnStream()
116 se_args[i]->shape_tree.buffers = in ExecuteAsyncOnStream()
119 se_args[i]->shape_tree.buffers[j] = se_buffers[j]; in ExecuteAsyncOnStream()
143 ApiConverter::Free(&se_args[i]->shape_tree.shape); in ExecuteAsyncOnStream()
146 delete[] se_args[i]->shape_tree.buffers; in ExecuteAsyncOnStream()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | BUILD | 120 "//tensorflow/compiler/xla:shape_tree",
191 "//tensorflow/compiler/xla:shape_tree",
|
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | c_api_decl.h | 256 XLA_MaybeOwningDeviceMemoryShapeTree shape_tree; member
|
| /external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
| D | BUILD | 83 "//tensorflow/compiler/xla:shape_tree",
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | BUILD | 190 "//tensorflow/compiler/xla:shape_tree",
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | BUILD | 328 "//tensorflow/compiler/xla:shape_tree",
|