Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/xla/service/
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()
Dbfloat16_propagation.cc55 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineFusionComputationPrecision()
103 inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in RevertIfFusionInternalBF16Changes()
180 const Shape& subshape, const ShapeIndex& index) { in DetermineWhileComputationsPrecision()
215 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineConditionalComputationsPrecision()
238 const Shape& subshape = ShapeUtil::GetSubshape(hlo.shape(), index); in AllUsersConsumeBF16() local
945 auto subshape = entry.first; in Run() local
979 Shape* subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index); in OutputTypeAfterChange() local
Dhlo_replication_analysis.cc314 inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ComputeHloReplicationOnComputation()
339 param->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ComputeHloReplication()
351 param->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ComputeHloReplication()
Dbfloat16_normalization.cc98 shape, [&](const Shape& subshape, const ShapeIndex& index) { in CountSubshapesWithMatchingType()
109 shape, [&](const Shape& subshape, const ShapeIndex& index) { in ShapeLeafCount()
178 hlo->mutable_shape(), [&](Shape* subshape, const xla::ShapeIndex& index) { in ChangeOutputTypeThenInsertConvertBack()
351 auto subshape = ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), {i}); in HandleMultipleOutputs() local
Dwhile_loop_invariant_code_motion.cc243 [&output_size, this](const Shape& subshape, in TryHoistingInvariantInstructionsFromWhileBody()
244 const ShapeIndex& /*index*/) { in TryHoistingInvariantInstructionsFromWhileBody()
Dlayout_assignment.cc353 const Shape& subshape, const ShapeIndex& index) -> Status { in SetInstructionLayout()
856 [&](const Shape& subshape, const ShapeIndex& shape_index) { in CheckParameterLayout()
1498 [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint()
1539 user->shape(), [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint()
1720 const Shape& subshape = ShapeUtil::GetSubshape(instruction->shape(), index); in InferArrayLayout() local
1851 [instruction, &constraints](Shape* subshape, const ShapeIndex& index) { in AssignLayouts()
2183 [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateComputationLayouts()
2475 return absl::c_all_of(shape.tuple_shapes(), [](const Shape& subshape) { in IsAtMostRank1()
Dinstruction_fusion.cc210 [&output_rank](const Shape& subshape, const ShapeIndex& shape_index) { in EffectivelyAtMostUnary()
304 shape, [&size](const Shape& subshape, const ShapeIndex& shape_index) { in ComputeGloballyUnfusible()
Dhlo_element_type_converter.cc70 Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); in GetConvertedTupleShape() local
Dgeneric_transfer_manager.cc74 [&](const Shape& subshape, const ShapeIndex& index) -> Status { in TransferLiteralFromDevice()
Dshaped_buffer.cc96 [this, &s](const Shape& subshape, const ShapeIndex& index) { in ToString()
Ddynamic_dimension_inference.cc439 [&](const Shape& subshape, ShapeIndex reduce_result_index) { in HandleReduce()
1120 [&](const Shape& subshape, ShapeIndex reduce_window_result_index) { in HandleReduceWindow()
1389 hlo->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in HandleConditional()
1413 const ShapeIndex& index) { in HandleConditional()
1727 Shape subshape = ShapeUtil::GetSubshape(inst->shape(), index); in SetDynamicSize() local
1822 const ShapeIndex& subindex) { in HasDynamicDimension()
/external/tensorflow/tensorflow/compiler/xla/
Dliteral.cc150 const Shape& subshape = shape.tuple_shapes(i); in SetPiece() local
726 Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); in Relayout() local
742 [this, &result](const Shape& subshape, const ShapeIndex& index) { in Relayout()
756 shape(), [&](const Shape& subshape, const ShapeIndex& index) { in ToBoundedDynamic()
773 &new_shape, [this](Shape* subshape, const ShapeIndex& index) { in ToStatic()
988 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString() local
1105 shape(), [&](const Shape& subshape, const ShapeIndex& index) { in Hash()
1185 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in TupleToStringHelper() local
1203 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in DenseArrayToStringHelper() local
1290 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in ToStringHelper() local
[all …]
Dshape_layout.cc56 const xla::ShapeIndex& index) { in MatchesLayoutInShape()
Dliteral_util.cc51 &result_shape, [](Shape* subshape, const ShapeIndex&) { in ConvertType()
64 [&](const Shape& subshape, const ShapeIndex& shape_index) { in ConvertType()
Dshape.cc107 for (const Shape& subshape : tuple_shapes_) { in is_static() local
Dshape_util.cc938 const Shape* subshape = &shape; in IndexIsValid() local
994 for (const Shape& subshape : shape.tuple_shapes()) { in GetLeafCount() local
1064 [&func](const Shape& subshape, const ShapeIndex& index) { in ForEachSubshape()
1077 [&func](Shape* subshape, const ShapeIndex& index) { in ForEachMutableSubshape()
1690 for (const Shape& subshape : shape.tuple_shapes()) { in Hash() local
1766 ForEachMutableSubshape(&s, [](Shape* subshape, const ShapeIndex& index) { in DeviceShapeToHostShape()
/external/tensorflow/tensorflow/core/ops/
Dmath_ops.cc1149 ShapeHandle subshape; in SegmentReductionShapeFn() local
1173 ShapeHandle subshape; in SparseSegmentReductionShapeFn() local
1197 ShapeHandle subshape; in SparseSegmentReductionGradShapeFn() local
1238 ShapeHandle subshape; in SparseSegmentReductionWithNumSegmentsShapeFn() local
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtracked_device_buffer_test.cc35 [&](const Shape& subshape, const ShapeIndex&) -> Status { in MakeArray()
Dutils.cc173 dst_shape, [&](Shape* subshape, const ShapeIndex& idx) { in DetermineArgumentLayoutsFromCompileOptions()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc54 xla::Shape subshape = in AllocateShapedBuffer() local
/external/tensorflow/tensorflow/compiler/xrt/
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
/external/tensorflow/tensorflow/compiler/xla/tests/
Dclient_library_test_base.cc332 &layout_shape, [&](Shape* subshape, const ShapeIndex& /*index*/) { in ComputeAndCompareLiteralWithStatus()
389 &layout_shape, [&](Shape* subshape, const ShapeIndex& /*index*/) { in ComputeAndCompareLiteralWithStatus()
622 [](Shape* subshape, const ShapeIndex&) { in MaybeConvertShapeToBfloat16()
/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_executable_interface.cc64 const xla::ShapeIndex& index) { in AllocateOutputMemoryWithInputReuse()
/external/tensorflow/tensorflow/compiler/xla/tools/
Dhlo_control_flow_flattening.cc85 Shape* subshape = shape->add_tuple_shapes(); in FlattenWhileLoop() local
/external/tensorflow/tensorflow/compiler/tf2xla/
Dshape_util.cc33 const xla::Shape& subshape = in PopulateInfeedLayoutVector() local

12