| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/ |
| H A D | shape_util.h | 73 struct ShapeIndex : public absl::InlinedVector<int64_t, 2> { struct
77 explicit ShapeIndex(ShapeIndexView view) in ShapeIndex() function
87 std::ostream& operator<<(std::ostream& out, const ShapeIndex& shape_index); argument
|
| H A D | shape_tree_test.cc | 66 int_tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestShapeConstructor()
75 [&num_nodes](const ShapeIndex& /*index*/, bool data) { in TestShapeConstructor()
92 tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestInitValueConstructor()
100 [&num_nodes](const ShapeIndex& /*index*/, int* data) { in TestInitValueConstructor()
108 tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestInitValueConstructor()
172 [&sum](const ShapeIndex& /*index*/, int data) { sum += data; }); in TEST_F()
184 [](const ShapeIndex& /*index*/, int* data) { *data = 0; }); in TEST_F()
391 t.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TEST_F()
|
| H A D | shape_tree.h | 217 absl::FunctionRef<void(const ShapeIndex&, const T&)> func) const { in ForEachElement()
224 absl::FunctionRef<void(const ShapeIndex&, T*)> func) { in ForEachMutableElement()
233 absl::FunctionRef<Status(const ShapeIndex&, const T&)> func) const { in ForEachElementWithStatus()
241 absl::FunctionRef<Status(const ShapeIndex&, T*)> func) { in ForEachMutableElementWithStatus()
|
| H A D | layout_util.cc | 287 shape, [](Shape* subshape, const ShapeIndex&) { in ClearTiles()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | tuple_points_to_analysis.cc | 54 [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsAmbiguous()
63 ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsDistinct()
83 [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) { in CreateFlattenedSet()
101 const ShapeIndex& index) const { in ContainsBufferAtIndex() argument
|
| H A D | while_loop_invariant_code_motion.cc | 240 const ShapeIndex& /*index*/) { in TryHoistingInvariantInstructionsFromWhileBody()
249 const ShapeIndex& /*index*/) { in TryHoistingInvariantInstructionsFromWhileBody()
|
| H A D | dynamic_dimension_inference.cc | 1308 hlo, [&](HloInstruction*, ShapeIndex /*index*/, int64_t dimension, in HandleDynamicSlice() argument
1332 [&](HloInstruction* /*operand*/, ShapeIndex /*index*/, int64_t dimension, in HandleDynamicUpdateSlice() argument
1333 int64_t operand_index, HloInstruction* dynamic_size) { in HandleDynamicUpdateSlice()
1368 HloInstruction* dynamic_size) { in HandleGather()
1458 HloInstruction* dynamic_size) -> Status { in HandleConditional()
1630 int64_t operand_index, HloInstruction* operand_dynamic_size) { in HandleScatter()
|
| H A D | optimize_input_output_buffer_alias_test.cc | 47 [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { in AliasCount()
|
| H A D | while_loop_simplifier.cc | 958 subshape, [&](const Shape& s, const ShapeIndex& /*index*/) { in UnflattenTupleInstr()
1020 [&](const Shape& s, const ShapeIndex& /*index*/) { in TryFlattenNestedTuples()
|
| H A D | hlo_input_output_alias_config.cc | 128 [&](const xla::ShapeIndex&, std::optional<Alias> alias) { in ParameterMustAlias()
|
| H A D | while_loop_expensive_invariant_code_motion.cc | 283 const ShapeIndex& /*index*/) { in TryHoistingInvariantInstructionsFromWhileBody()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/ |
| H A D | client_library_test_base.cc | 332 &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()
|
| H A D | literal_test_util.cc | 53 const ShapeIndex& /*shape_index*/, in OnMiscompare()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| H A D | cpu_shape_verifier.cc | 22 hlo->shape(), [&](const Shape& shape, const ShapeIndex&) { in Preprocess()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | gpu_shape_verifier.cc | 22 hlo->shape(), [&](const Shape& shape, const ShapeIndex&) { in Preprocess()
|
| H A D | gpu_executable.cc | 889 absl::flat_hash_map<ShapeIndex, GpuExecutable::OutputInfo>* output_info, in SetUpMlirAllocation()
1011 absl::flat_hash_map<ShapeIndex, OutputInfo> output_info, in GpuExecutable()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| H A D | tracked_device_buffer_test.cc | 35 [&](const Shape& subshape, const ShapeIndex&) -> Status { in MakeArray()
|
| H A D | tracked_device_buffer.cc | 109 shaped_buffer->on_device_shape(), [&](const Shape&, const ShapeIndex&) { in FromScopedShapedBuffer()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_executable_info.proto | 39 message ShapeIndex { message
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tools/ |
| H A D | run_hlo_module.cc | 82 const ShapeIndex& /*shape_index*/, in OnMiscompare()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/ |
| H A D | ops.cc | 200 CustomCallApiVersion api_version) -> XlaOp { in BuildOpsSubmodule()
|
| H A D | outfeed_receiver.cc | 430 &shape_with_layout, [](Shape* subshape, const ShapeIndex&) { in AddOutfeedToBuilder()
|
| H A D | xla.cc | 163 &shape_with_layout, [](Shape* subshape, const ShapeIndex&) { in PYBIND11_MODULE()
|
| H A D | xla_client.py | 233 ShapeIndex = _xla.ShapeIndex variable
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | xla_builder.cc | 1970 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCall()
2018 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCallInternal()
2079 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCall()
4424 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCall()
4439 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCallWithComputation()
4454 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCallWithLayout()
4469 absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> in CustomCallWithConvDnums()
|