Searched refs:sub_shape (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding.cc | 450 const Shape* sub_shape = &shape; in GetSubSharding() local
454 ShapeUtil::GetLeafCount(ShapeUtil::GetSubshape(*sub_shape, {i})); in GetSubSharding()
456 sub_shape = &ShapeUtil::GetSubshape(*sub_shape, {idx}); in GetSubSharding()
458 if (sub_shape->IsTuple()) { in GetSubSharding()
461 begin_it, begin_it + ShapeUtil::GetLeafCount(*sub_shape)); in GetSubSharding()
462 return HloSharding::Tuple(*sub_shape, sub_shardings); in GetSubSharding()
|
| D | hlo_liveness_analysis.cc | 99 [&](const Shape& sub_shape, const ShapeIndex& shape_index) { in MarkLiveAtAllIndices() argument
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | shape_util.cc | 836 ForEachSubshape(shape, [&](const Shape& sub_shape, const ShapeIndex& index) { in GetLeafShapes() argument
838 leaves.emplace_back(index, sub_shape); in GetLeafShapes()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | ir_emitter_unnested.cc | 304 ShapeUtil::ForEachSubshape(s, [&](const Shape& sub_shape, in GetIndexTypeForKernel()
306 if (sub_shape.IsArray() && !IsInt32(ShapeUtil::ElementsIn(sub_shape))) { in GetIndexTypeForKernel()
|