Searched refs:GetLeafCount (Results 1 – 6 of 6) sorted by relevance
205 return ShapeUtil::IsEmptyTuple(shape) ? 1 : ShapeUtil::GetLeafCount(shape); in RequiredLeaves()454 ShapeUtil::GetLeafCount(ShapeUtil::GetSubshape(*sub_shape, {i})); in GetSubSharding()461 begin_it, begin_it + ShapeUtil::GetLeafCount(*sub_shape)); in GetSubSharding()
238 ShapeUtil::GetLeafCount(instruction->shape()) == in ApplyShardingFromUsers()
872 CHECK_EQ(ShapeUtil::GetLeafCount(shape()), in set_parameter_replicated_at_leaf_buffers()
1726 int leaf_count = ShapeUtil::GetLeafCount(instruction->shape()); in ParseInstructionRhs()
822 /* static */ int64 ShapeUtil::GetLeafCount(const Shape& shape) { in GetLeafCount() function in xla::ShapeUtil828 count += GetLeafCount(subshape); in GetLeafCount()
502 static int64 GetLeafCount(const Shape& shape);