Home
last modified time | relevance | path

Searched refs:GetMutableSubshape (Results 1 – 25 of 25) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dbfloat16_conversion_folding.cc243 ShapeUtil::GetMutableSubshape(crs->mutable_shape(), {i}) in HandleAllReduce()
246 ShapeUtil::GetMutableSubshape(crs->mutable_shape(), {i})); in HandleAllReduce()
H A Dbfloat16_propagation.cc931 ShapeUtil::GetMutableSubshape( in Run()
986 Shape* subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index); in OutputTypeAfterChange()
1009 entry.emplace(ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index), in AddToOrRemoveFromBF16ChangeSet()
1018 ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index)); in AddToOrRemoveFromBF16ChangeSet()
H A Dmemory_space_propagation.cc73 Shape* shape = ShapeUtil::GetMutableSubshape(instruction->mutable_shape(), in Propagate()
H A Dtransfer_manager.cc214 *ShapeUtil::GetMutableSubshape(device_shape, index); in ReadDynamicShapes()
H A Dbfloat16_normalization.cc350 auto subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), {i}); in HandleMultipleOutputs()
H A Dlayout_assignment_test.cc276 *ShapeUtil::GetMutableSubshape(&result_shape, /*index=*/{0, 0}) = in TEST_F()
278 *ShapeUtil::GetMutableSubshape(&result_shape, /*index=*/{1, 0}) = in TEST_F()
H A Dhlo_instruction_test.cc843 *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(), {0}) in TEST_F()
845 *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(), {1}) in TEST_F()
H A Dlayout_assignment.cc1989 Shape* buffer_subshape = ShapeUtil::GetMutableSubshape( in AssignLayouts()
2352 Shape* shape = ShapeUtil::GetMutableSubshape( in PropagateMemorySpace()
H A Dbuffer_assignment.cc247 Shape* shape = ShapeUtil::GetMutableSubshape( in AddAssignment()
H A Dmemory_space_assignment_test.cc3104 *ShapeUtil::GetMutableSubshape(&tuple_shape, {0})->mutable_layout() = in TEST_P()
3109 *ShapeUtil::GetMutableSubshape(&tuple_shape, {1})->mutable_layout() = in TEST_P()
3114 *ShapeUtil::GetMutableSubshape(&tuple_shape, {2})->mutable_layout() = in TEST_P()
H A Dmemory_space_assignment.cc1503 Shape* shape = ShapeUtil::GetMutableSubshape( in Finish()
3908 Shape* shape = ShapeUtil::GetMutableSubshape( in ExportAndColorBuffers()
H A Dpattern_matcher.h955 return ShapeUtil::GetMutableSubshape(shape, index_);
H A Dhlo_instructions.cc1475 ShapeUtil::GetMutableSubshape(mutable_shape(), shape_index); in RelayoutConstant()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dshape_test.cc127 ShapeUtil::GetMutableSubshape(&dynamic_tuple, {2}) in TEST_F()
140 ShapeUtil::GetMutableSubshape(&dynamic_tuple, {2}) in TEST_F()
H A Dshape_layout.cc101 *ShapeUtil::GetMutableSubshape(&shape_, shape_index)->mutable_layout() = in ResetLayout()
H A Dshape_util_test.cc458 array_shape, *ShapeUtil::GetMutableSubshape(&array_shape, {}))); in TEST()
557 EXPECT_EQ(subshape, ShapeUtil::GetMutableSubshape(&shape, index)); in TEST()
H A Dshape_util.h456 static Shape* GetMutableSubshape(Shape* shape, ShapeIndexView index);
H A Dliteral.cc677 ShapeUtil::GetMutableSubshape(mutable_shape_do_not_use(), shape_index); in SetDynamicSize()
853 Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); in Relayout()
H A Dshape_util.cc1013 /* static */ Shape* ShapeUtil::GetMutableSubshape(Shape* shape, in GetMutableSubshape() function in xla::ShapeUtil
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/
H A Dxrt_state.cc439 *xla::ShapeUtil::GetMutableSubshape(host_shape, index) = in ExpandTreeOfTuples()
441 *xla::ShapeUtil::GetMutableSubshape(device_shape, index) = in ExpandTreeOfTuples()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_transfer_manager.cc108 *ShapeUtil::GetMutableSubshape(device_shape, index); in ReadDynamicShapes()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/
H A Dxrt_state_ops.h81 *xla::ShapeUtil::GetMutableSubshape(shape, {}) = in ParseTupleNode()
86 xla::ShapeUtil::GetMutableSubshape(shape, {i}), rm)); in ParseTupleNode()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/
H A Dcpu_xfeed.cc289 *ShapeUtil::GetMutableSubshape(device_shape, index); in ReadDynamicShapesOnCpu()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dinfeed_ops.cc137 xla::ShapeUtil::GetMutableSubshape(output_shape, {i}); in GetInfeedShapeWithLayout()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Ddot_handler.cc3890 *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(), in SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions()
4165 *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(), {2}) = in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions()