Home
last modified time | relevance | path

Searched refs:scatter_op (Results 1 – 16 of 16) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Dscatter_functor_gpu.cu.cc31 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ASSIGN); \
32 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ADD); \
33 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::SUB); \
34 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MUL); \
35 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::DIV); \
36 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MIN); \
37 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MAX);
46 DEFINE_GPU_SPECS_OP(bool, int32, scatter_op::UpdateOp::ASSIGN);
47 DEFINE_GPU_SPECS_OP(bool, int64, scatter_op::UpdateOp::ASSIGN);
Dscatter_op_gpu.cu.cc32 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ASSIGN); \
33 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ADD); \
34 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::SUB); \
35 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MUL); \
36 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::DIV); \
37 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MIN); \
38 DEFINE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MAX);
Dscatter_functor.h37 namespace scatter_op {
43 template <scatter_op::UpdateOp Op>
46 struct Assign<scatter_op::UpdateOp::ASSIGN> {
57 struct Assign<scatter_op::UpdateOp::ADD> {
68 struct Assign<scatter_op::UpdateOp::SUB> {
79 struct Assign<scatter_op::UpdateOp::MUL> {
90 struct Assign<scatter_op::UpdateOp::DIV> {
101 struct Assign<scatter_op::UpdateOp::MIN> {
114 struct Assign<scatter_op::UpdateOp::MAX> {
130 template <typename Device, typename T, typename Index, scatter_op::UpdateOp op>
[all …]
Dscatter_functor.cc46 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ASSIGN); \
47 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ADD); \
48 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::SUB); \
49 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MUL); \
50 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::DIV); \
51 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MIN); \
52 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MAX);
Dscatter_functor_gpu.cu.h34 template <typename T, scatter_op::UpdateOp op>
38 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::ASSIGN> {
43 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::ADD> {
50 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::SUB> {
57 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::MUL> {
64 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::DIV> {
71 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::MIN> {
78 struct ScatterOpKernelBody<T, scatter_op::UpdateOp::MAX> {
84 template <typename T, typename Index, scatter_op::UpdateOp op>
105 template <typename T, typename Index, scatter_op::UpdateOp op>
[all …]
Dscatter_op.cc66 template <typename Device, typename T, typename Index, scatter_op::UpdateOp op>
161 REGISTER_SCATTER_KERNEL(type, dev, "ScatterAdd", scatter_op::UpdateOp::ADD); \
162 REGISTER_SCATTER_KERNEL(type, dev, "ScatterDiv", scatter_op::UpdateOp::DIV); \
163 REGISTER_SCATTER_KERNEL(type, dev, "ScatterMul", scatter_op::UpdateOp::MUL); \
164 REGISTER_SCATTER_KERNEL(type, dev, "ScatterSub", scatter_op::UpdateOp::SUB);
167 REGISTER_SCATTER_KERNEL(type, dev, "ScatterMin", scatter_op::UpdateOp::MIN); \
168 REGISTER_SCATTER_KERNEL(type, dev, "ScatterMax", scatter_op::UpdateOp::MAX);
172 scatter_op::UpdateOp::ASSIGN);
Dresource_variable_ops.cc837 template <typename Device, typename T, typename Index, scatter_op::UpdateOp op>
954 scatter_op::UpdateOp::ADD); \
956 scatter_op::UpdateOp::SUB); \
958 scatter_op::UpdateOp::MUL); \
960 scatter_op::UpdateOp::DIV); \
962 scatter_op::UpdateOp::ASSIGN);
965 scatter_op::UpdateOp::MIN); \
967 scatter_op::UpdateOp::MAX);
978 scatter_op::UpdateOp::ASSIGN);
980 scatter_op::UpdateOp::ASSIGN);
[all …]
Dsparse_tensor_dense_add_op.cc97 scatter_op::UpdateOp::ADD>()( \ in Compute()
125 struct ScatterNdFunctor<CPUDevice, T, Index, NDIMS, scatter_op::UpdateOp::ADD> {
Dsparse_tensor_dense_add_op.h31 scatter_op::UpdateOp op>
DBUILD4826 ":scatter_op",
4855 name = "scatter_op",
4856 prefix = "scatter_op",
4938 ":scatter_op",
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc1092 mhlo::ScatterOp scatter_op, ArrayRef<Value> args, in matchAndRewrite() argument
1094 Value operand = scatter_op.operand(); in matchAndRewrite()
1095 Value indices = scatter_op.scatter_indices(); in matchAndRewrite()
1096 Value updates = scatter_op.updates(); in matchAndRewrite()
1109 scatter_op.scatter_dimension_numbers().index_vector_dim().getInt(); in matchAndRewrite()
1110 if (failed(NormalizeIndexVector(scatter_op, indices, indices_type, in matchAndRewrite()
1119 scatter_op.scatter_dimension_numbers().inserted_window_dims(); in matchAndRewrite()
1121 scatter_op.scatter_dimension_numbers().scatter_dims_to_operand_dims(); in matchAndRewrite()
1128 scatter_op, in matchAndRewrite()
1136 scatter_op.scatter_dimension_numbers().update_window_dims(); in matchAndRewrite()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc442 auto scatter_op = func_builder->create<mlir::mhlo::ScatterOp>( in ImportInstructionImpl() local
445 &scatter_op.update_computation())); in ImportInstructionImpl()
446 return scatter_op.getOperation(); in ImportInstructionImpl()
/external/tensorflow/tensorflow/python/kernel_tests/
Dembedding_ops_test.py56 def _TestCase(self, shape, indices, scatter_op=state_ops.scatter_add): argument
77 p2 = scatter_op(p, ind, vals, name="updated_p")
84 if scatter_op == state_ops.scatter_add:
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc2502 auto scatter_op = mlir::cast<mlir::lmhlo::ScatterOp>(mlir_input.op); in EmitScatterFromMlir() local
2505 GetAllocationSliceForMlir(scatter_op.operand())); in EmitScatterFromMlir()
2507 GetAllocationSliceForMlir(scatter_op.output())); in EmitScatterFromMlir()
2516 ShapeUtil::ByteSizeOf(TypeToShape(scatter_op.output().getType())))); in EmitScatterFromMlir()
2526 BuildKernelThunkForMlir(scatter_op, scatter_op.getOperands().drop_front(), in EmitScatterFromMlir()
2536 return GetIndexTypeForKernelFromMlir(scatter_op, launch_size, &b_); in EmitScatterFromMlir()
2540 thunks.back().get(), scatter_op, output, in EmitScatterFromMlir()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dhlo_text_to_lhlo_no_opt.hlotxt42 ROOT scatter_op = s32[3,3] scatter(operand, indices, updates),
/external/llvm-project/mlir/test/Conversion/VectorToLLVM/
Dvector-to-llvm.mlir1102 func @scatter_op(%arg0: memref<?xf32>, %arg1: vector<3xi32>, %arg2: vector<3xi1>, %arg3: vector<3xf…
1107 // CHECK-LABEL: func @scatter_op