Home
last modified time | relevance | path

Searched defs:subshape (Results 1 – 25 of 56) sorted by relevance

123

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dliteral.cc233 const Shape& subshape = shape.tuple_shapes(i); in mutable_shape_do_not_use() local
251 for (const Shape& subshape : shape.tuple_shapes()) { in SetPiece() local
676 Shape* subshape = in SetDynamicSize() local
853 Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); in Relayout() local
869 [this, &result](const Shape& subshape, const ShapeIndex& index) { in Relayout()
883 shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ToBoundedDynamic()
900 &new_shape, [this](Shape* subshape, const ShapeIndex& index) { in ToStatic()
1128 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString() local
1304 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in TupleToStringHelper() local
1325 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in DenseArrayToStringHelper() local
[all …]
H A Dshape_util.cc977 const Shape* subshape = &shape; in IndexIsValid() local
1033 for (const Shape& subshape : shape.tuple_shapes()) { in GetLeafCount() local
1103 [&func](const Shape& subshape, const ShapeIndex& index) { in ForEachSubshape()
1116 [&func](Shape* subshape, const ShapeIndex& index) { in ForEachMutableSubshape()
1729 ForEachMutableSubshape(&s, [](Shape* subshape, const ShapeIndex& index) { in DeviceShapeToHostShape()
H A Dliteral_util.cc50 &result_shape, [](Shape* subshape, const ShapeIndex&) { in ConvertType()
63 [&](const Shape& subshape, const ShapeIndex& shape_index) { in ConvertType()
H A Dshape_util_test.cc518 shape, [&calls, &shape](const Shape& subshape, const ShapeIndex& index) { in TEST()
533 shape, [&calls, &shape](const Shape& subshape, const ShapeIndex& index) { in TEST()
555 &shape, [&calls, &shape](const Shape* subshape, const ShapeIndex& index) { in TEST()
H A Dshape_layout.cc56 const xla::ShapeIndex& index) { in MatchesLayoutInShape()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dbfloat16_propagation.cc55 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineFusionComputationPrecision()
103 const ShapeIndex& index) { in RevertIfFusionInternalBF16Changes()
176 const Shape& subshape, const ShapeIndex& index) { in DetermineWhileComputationsPrecision()
211 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineConditionalComputationsPrecision()
234 const Shape& subshape = ShapeUtil::GetSubshape(hlo.shape(), index); in AllUsersConsumeBF16() local
951 auto subshape = entry.first; in Run() local
986 Shape* subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index); in OutputTypeAfterChange() local
H A Dlayout_assignment.cc378 const Shape& subshape, const ShapeIndex& index) -> Status { in SetInstructionLayout()
416 const Shape& subshape, const ShapeIndex& index) -> Status { in SetInstructionLayout()
931 [&](const Shape& subshape, const ShapeIndex& shape_index) { in CheckParameterLayout()
1645 [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint()
1680 user->shape(), [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint()
1999 [instruction, this](Shape* subshape, const ShapeIndex& index) { in AssignLayouts()
2055 const xla::ShapeIndex& index) { in AssignLayouts()
2087 const Shape& subshape, const ShapeIndex& index) { in CalculateComputationLayout()
2372 [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateComputationLayouts()
2672 return absl::c_all_of(shape.tuple_shapes(), [](const Shape& subshape) { in IsAtMostRank1()
H A Dbfloat16_normalization.cc97 shape, [&](const Shape& subshape, const ShapeIndex& index) { in CountSubshapesWithMatchingType()
108 shape, [&](const Shape& subshape, const ShapeIndex& index) { in ShapeLeafCount()
177 hlo->mutable_shape(), [&](Shape* subshape, const xla::ShapeIndex& index) { in ChangeOutputTypeThenInsertConvertBack()
350 auto subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), {i}); in HandleMultipleOutputs() local
H A Doptimize_input_output_buffer_alias.cc57 input_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
67 output_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
H A Ddynamic_dimension_inference.cc470 [&](const Shape& subshape, ShapeIndex reduce_result_index) { in HandleReduce()
1258 [&](const Shape& subshape, ShapeIndex reduce_window_result_index) { in HandleReduceWindow()
1527 hlo->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in HandleConditional()
1551 hlo->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in HandleConditional()
1707 [&](const Shape& subshape, const ShapeIndex& index) { in HandleWhile()
1918 Shape subshape = ShapeUtil::GetSubshape(inst->shape(), index); in SetDynamicSize() local
2020 const ShapeIndex& subindex) { in HasDynamicDimension()
H A Dwhile_loop_invariant_code_motion.cc248 [&output_size, this](const Shape& subshape, in TryHoistingInvariantInstructionsFromWhileBody()
249 const ShapeIndex& /*index*/) { in TryHoistingInvariantInstructionsFromWhileBody()
H A Dhlo_replication_analysis.cc382 inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ComputeHloReplicationOnComputation()
408 param->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ComputeHloReplication()
H A Dhlo_dataflow_analysis.cc612 operand->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in UpdateAsyncStartValueSet()
633 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in UpdateAsyncStartValueSet()
662 [&](const Shape& subshape, const ShapeIndex& index) { in UpdateAsyncUpdateValueSet()
697 [&](const Shape& subshape, const ShapeIndex& index) { in UpdateAsyncDoneValueSet()
H A Dinstruction_fusion.cc215 [&output_rank](const Shape& subshape, const ShapeIndex& shape_index) { in EffectivelyAtMostUnary()
311 shape, [&size](const Shape& subshape, const ShapeIndex& shape_index) { in ComputeGloballyUnfusible()
H A Dhlo_cost_analysis.cc792 crs->shape(), [&](const Shape& subshape, const ShapeIndex&) { in HandleAllReduce()
891 [this, fusion](const Shape& subshape, const ShapeIndex& shape_index) { in HandleFusion()
938 const Shape& subshape = shape.tuple_shapes(i); in HandleFusion() local
H A Dcompile_only_service.cc104 const xla::ShapeIndex& index) { in CompileAheadOfTime()
H A Dhlo_module_util.cc157 Shape* subshape, const ShapeIndex& index) { in UpdateEntryComputationLayout()
H A Dtuple_util.cc94 const Shape& subshape = tuple_shape.tuple_shapes(i); in ReplaceTupleWith() local
H A Dhlo_element_type_converter.cc72 Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); in GetConvertedTupleShape() local
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/
H A Dmath_ops.cc1161 ShapeHandle subshape; in SegmentReductionShapeFn() local
1185 ShapeHandle subshape; in SparseSegmentReductionShapeFn() local
1209 ShapeHandle subshape; in SparseSegmentReductionGradShapeFn() local
1250 ShapeHandle subshape; in SparseSegmentReductionWithNumSegmentsShapeFn() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dvalue_inference.cc392 Shape subshape = in IsInstructionOverLimit() local
477 Shape subshape = in AnalyzeConstantValueFallback() local
645 Shape subshape = in AnalyzeUpperBound() local
815 Shape subshape = in AnalyzeLowerBound() local
905 Shape subshape = in AnalyzeConstant() local
1007 Shape subshape = in AnalyzeIsDynamic() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/
H A Dxrt_state.cc99 const xla::Shape& subshape = in AllocateScopedShapedBuffer() local
364 XRTTupleAllocation* parent, const xla::ShapeIndex& subshape, in MakeSubBuffer()
485 const xla::Shape& subshape = in MakeTuple() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dclient_library_test_base.cc332 &layout_shape, [&](Shape* subshape, const ShapeIndex& /*index*/) { in ComputeAndCompareLiteralWithStatus()
390 &layout_shape, [&](Shape* subshape, const ShapeIndex& /*index*/) { in ComputeAndCompareLiteralWithStatus()
624 [](Shape* subshape, const ShapeIndex&) { in MaybeConvertShapeToBfloat16()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dtracked_device_buffer_test.cc35 [&](const Shape& subshape, const ShapeIndex&) -> Status { in MakeArray()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dxla_tensor.cc54 xla::Shape subshape = in AllocateShapedBuffer() local

123