experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fixes a previously unhandled case of a tuple with zero leaf nodes in

Browse Source 
hlo_sharding.cc.

PiperOrigin-RevId: 306997392
Change-Id: I86d006f8004e9ebb93baf1cd79bac9443b8dabbc
This commit is contained in:
A. Unique TensorFlower 2020-04-17 00:26:45 -07:00 committed by TensorFlower Gardener
parent 5dd47ad28a
commit b3cc88cdbc
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/compiler/xla/service/hlo_sharding.cc
View File

@ -199,10 +199,12 @@ std::vector<int64> HloSharding::TileLimitForDevice(const Shape& shape,
} }
int64 HloSharding::RequiredLeaves(const Shape& shape) { int64 HloSharding::RequiredLeaves(const Shape& shape) {
// Empty tuples have no leaf nodes as far as ShapeUtil and ShapeTree are // Empty tuples (with arbitrary nesting) have no leaf nodes as far as
// concerned, but they do have a single tuple_elements_ entry since we want // ShapeUtil and ShapeTree are concerned, but they do have a single
// to allow empty tuple results to have sharding. // tuple_elements_ entry since we want to allow empty tuple results to
return ShapeUtil::IsEmptyTuple(shape) ? 1 : ShapeUtil::GetLeafCount(shape); // have sharding.
const int64 leaf_count = ShapeUtil::GetLeafCount(shape);
return (leaf_count == 0) ? 1 : leaf_count;
} }
Status HloSharding::CheckLeafCount(const Shape& shape) const { Status HloSharding::CheckLeafCount(const Shape& shape) const {