Searched refs:GetLeafCount (Results 1 – 12 of 12) sorted by relevance
473 const int64_t leaf_count = ShapeUtil::GetLeafCount(shape); in RequiredLeaves()478 int64_t leaf_count = ShapeUtil::GetLeafCount(shape); in CheckLeafCount()551 if (ShapeUtil::GetLeafCount(shape) == 0 && tuple_elements_.empty()) { in ValidateTuple()813 ShapeUtil::GetLeafCount(ShapeUtil::GetSubshape(*sub_shape, {i})); in GetSubSharding()820 begin_it, begin_it + ShapeUtil::GetLeafCount(*sub_shape)); in GetSubSharding()
242 ShapeUtil::GetLeafCount(instruction->shape()) == in ApplyShardingFromUsers()
1475 ShapeUtil::GetLeafCount(instruction.shape().tuple_shapes(i)); in GetShardingFromUser()1835 for (int64_t i = 0, e = ShapeUtil::GetLeafCount(operand->shape()); in InferShardingFromOperands()1851 sub_sharding_index += ShapeUtil::GetLeafCount(operand->shape()); in InferShardingFromOperands()
1398 CHECK_EQ(ShapeUtil::GetLeafCount(shape()), in set_parameter_replicated_at_leaf_buffers()1406 CHECK_EQ(ShapeUtil::GetLeafCount(shape()), in set_parameter_replicated_at_leaf_buffers()
1210 int leaf_count = ShapeUtil::GetLeafCount(instruction->shape()); in ParseInstructionRhs()
387 xla::ShapeUtil::GetLeafCount(elem_shapes[i])); in BuildComputation()1096 xla::ShapeUtil::GetLeafCount(arg_shapes[i]), in BuildArguments()1129 xla::ShapeUtil::GetLeafCount((*input_shapes)[i]), in BuildArguments()
463 static int64_t GetLeafCount(const Shape& shape);
130 EXPECT_EQ(ShapeUtil::GetLeafCount(shape), shape_tree.leaf_count()); in TEST_F()
1028 /* static */ int64_t ShapeUtil::GetLeafCount(const Shape& shape) { in GetLeafCount() function in xla::ShapeUtil1034 count += GetLeafCount(subshape); in GetLeafCount()
339 [](const Shape& shape) { return ShapeUtil::GetLeafCount(shape); }) in BuildXlaCompilerSubmodule()
2348 for (int i = 0, e = xla::ShapeUtil::GetLeafCount(arg_shape); i < e; ++i) in SetEntryTupleShapesAndLeafReplication()2483 xla::ShapeUtil::GetLeafCount(shape))); in LowerBasicBlockAsFunction()
3249 if (ShapeUtil::GetLeafCount(shape) == 0) { in HandleInfeed()