Home
last modified time | relevance | path

Searched refs:GetLeafCount (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding.cc205 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()
Dhlo_sharding_metadata.cc238 ShapeUtil::GetLeafCount(instruction->shape()) == in ApplyShardingFromUsers()
Dhlo_instructions.h938 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()
Dhlo_parser.cc1836 int leaf_count = ShapeUtil::GetLeafCount(instruction->shape()); in ParseInstructionRhs()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler.cc372 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()
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util.cc822 /* static */ int64 ShapeUtil::GetLeafCount(const Shape& shape) { in GetLeafCount() function in xla::ShapeUtil
828 count += GetLeafCount(subshape); in GetLeafCount()
Dshape_util.h515 static int64 GetLeafCount(const Shape& shape);
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc983 for (int i = 0; i < xla::ShapeUtil::GetLeafCount(arg_shapes.back()); in LowerBasicBlockAsFunction()
1008 xla::ShapeUtil::GetLeafCount(shape))); in LowerBasicBlockAsFunction()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla.cc411 [](const Shape& shape) { return ShapeUtil::GetLeafCount(shape); }) in PYBIND11_MODULE()