| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| H A D | gather_scatter_ops.cc | 81 XlaCompiler::CompilationResult update_computation; in Compile() local
85 &update_computation)); in Compile()
89 ctx->Input("updates"), *update_computation.computation, in Compile()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/python/ |
| H A D | xla.py | 582 def scatter(operand, scatter_indices, updates, update_computation, argument
588 update_computation=update_computation,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ir/ |
| H A D | mlir_hlo_builder.cc | 396 absl::Span<const XlaOp> updates, const XlaComputation& update_computation, in ScatterInternal() argument
413 ImportComputation(update_computation.proto(), &op.update_computation())); in ScatterInternal()
|
| H A D | mlir_hlo_builder.h | 196 absl::Span<const XlaOp> updates, const XlaComputation& update_computation,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | xla_builder.h | 884 const XlaComputation& update_computation,
889 const XlaComputation& update_computation,
895 absl::Span<const XlaOp> updates, const XlaComputation& update_computation,
1475 const XlaComputation& update_computation,
1480 const XlaComputation& update_computation,
2601 const XlaComputation& update_computation,
2606 const XlaComputation& update_computation,
|
| H A D | xla_builder.cc | 2428 const XlaComputation& update_computation, in Scatter() argument
2432 absl::MakeConstSpan(&updates, 1), update_computation, in Scatter()
2438 const XlaComputation& update_computation, in Scatter() argument
2464 update_computation.GetProgramShape()); in Scatter()
2469 update_computation, dimension_numbers, in Scatter()
2476 absl::Span<const XlaOp> updates, const XlaComputation& update_computation, in ScatterInternal() argument
2486 AddCalledComputation(update_computation, &instr); in ScatterInternal()
4907 const XlaOp updates, const XlaComputation& update_computation, in Scatter() argument
4911 update_computation, dimension_numbers, in Scatter()
4917 const XlaComputation& update_computation, in Scatter() argument
[all …]
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/ |
| H A D | ops.pyi | 330 update_computation: XlaComputation,
339 update_computation: XlaComputation,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/thlo/IR/ |
| H A D | thlo_ops.td | 110 - update_computation is sum.
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/lib/Dialect/mhlo/transforms/ |
| H A D | legalize_mhlo_to_thlo.cc | 239 if (matchUpdateComputation(op.update_computation()).failed()) in matchAndRewrite()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ |
| H A D | mlir_hlo_to_hlo.cc | 1473 xla::XlaComputation update_computation; in ExportXlaOp() local
1474 if (failed(ctx.converter->LowerRegionAsComputation(&op.update_computation(), in ExportXlaOp()
1475 &update_computation))) { in ExportXlaOp()
1491 update_computation, dimension_numbers, in ExportXlaOp()
|
| H A D | hlo_function_importer.cc | 813 &scatter_op.update_computation(), in ImportInstructionImpl()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | ir_emitter_unnested.h | 487 const HloComputation* update_computation; member
|
| H A D | ir_emitter_unnested.cc | 2288 const HloComputation* update_computation, in EmitScatter()
2299 desc.update_computation = update_computation; in EmitScatter()
2424 *desc.update_computation, output_address, input_address, in EmitScatter()
2427 return EmitCallToNestedComputation(*desc.update_computation, in EmitScatter()
5279 desc.update_computation = root->called_computations()[0]; in EmitScatter()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | hlo_instruction.h | 1068 HloComputation* update_computation,
1075 HloComputation* update_computation,
|
| H A D | hlo_instruction.cc | 1938 HloComputation* update_computation, in CreateScatter() argument
1942 shape, {operand, scatter_indices, updates}, update_computation, in CreateScatter()
1949 HloComputation* update_computation, in CreateScatter() argument
1958 shape, args, update_computation, scatter_dim_numbers, indices_are_sorted, in CreateScatter()
|
| H A D | hlo_instruction_test.cc | 1678 auto* update_computation = in TEST_F() local
1684 update_computation, in TEST_F()
|
| H A D | hlo_parser.cc | 2861 optional<HloComputation*> update_computation; in CreateInstruction() local
2863 &update_computation}; in CreateInstruction()
2896 arg_shapes, update_computation.value()->ComputeProgramShape(), in CreateInstruction()
2905 operand_span.last(input_count), *update_computation, dim_numbers, in CreateInstruction()
|
| H A D | hlo_instructions.cc | 3142 HloComputation* update_computation, in HloScatterInstruction() argument
3152 AppendComputation(update_computation); in HloScatterInstruction()
|
| H A D | hlo_instructions.h | 2153 HloComputation* update_computation,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| H A D | operation_semantics.md | 2511 `update_computation`.
2516 <b> `scatter(operands..., scatter_indices, updates..., update_computation,
2525 `update_computation` | `XlaComputation` | Computation to be used for combining the e…
2617 =`update_computation`(`output`[`0`][`O`], ...,
2624 Note that the first parameter that is passed into the `update_computation` will
2627 specifically for cases when the `update_computation` is _not commutative_.
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/lhlo/IR/ |
| H A D | lhlo_ops.td | 958 updated with the values in `updates` using `update_computation`.
972 let regions = (region SizedRegion<1>:$update_computation);
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/ |
| H A D | import.hlotxt | 1095 %update_computation {
1105 …ow_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=1, to_apply=%update_computation
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/dialect/ |
| H A D | StablehloOps.td | 1874 updated with the values in `updates` using `update_computation`.
1887 let regions = (region SizedRegion<1>:$update_computation);
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| H A D | hlo_ops.td | 2045 updated with the values in `updates` using `update_computation`.
2058 let regions = (region SizedRegion<1>:$update_computation);
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/lib/Dialect/mhlo/IR/ |
| H A D | hlo_ops.cc | 7968 Block& block = update_computation().front(); in verify()
8156 llvm::hasSingleElement(update_computation().front())) { in fold()
8252 auto newValue = evaluateMhloRegion(update_computation(), {lhs, rhs}); in fold()
|