Searched refs:GetLeafCount (Results 1 – 9 of 9) 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()
938 CHECK_EQ(ShapeUtil::GetLeafCount(shape()), in set_parameter_replicated_at_leaf_buffers()946 CHECK_EQ(ShapeUtil::GetLeafCount(shape()), in set_parameter_replicated_at_leaf_buffers()
1836 int leaf_count = ShapeUtil::GetLeafCount(instruction->shape()); in ParseInstructionRhs()
372 xla::ShapeUtil::GetLeafCount(elem_shapes[i])); in BuildComputation()1033 xla::ShapeUtil::GetLeafCount(arg_shapes[i]), in BuildArguments()1075 xla::ShapeUtil::GetLeafCount((*input_shapes)[i]), in BuildArguments()
822 /* static */ int64 ShapeUtil::GetLeafCount(const Shape& shape) { in GetLeafCount() function in xla::ShapeUtil828 count += GetLeafCount(subshape); in GetLeafCount()
515 static int64 GetLeafCount(const Shape& shape);
983 for (int i = 0; i < xla::ShapeUtil::GetLeafCount(arg_shapes.back()); in LowerBasicBlockAsFunction()1008 xla::ShapeUtil::GetLeafCount(shape))); in LowerBasicBlockAsFunction()
411 [](const Shape& shape) { return ShapeUtil::GetLeafCount(shape); }) in PYBIND11_MODULE()