Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/tf2xla/ops/
Dxla_ops.cc613 std::vector<int64_t> dimensions_to_reduce; in __anon5f0f96000502() local
615 c->GetAttr("dimensions_to_reduce", &dimensions_to_reduce)); in __anon5f0f96000502()
616 std::set<int64_t> dims_set(dimensions_to_reduce.begin(), in __anon5f0f96000502()
617 dimensions_to_reduce.end()); in __anon5f0f96000502()
621 const int dimensions_to_reduce_size = dimensions_to_reduce.size(); in __anon5f0f96000502()
623 dims_set.size() != dimensions_to_reduce.size() || in __anon5f0f96000502()
624 !absl::c_all_of(dimensions_to_reduce, dim_in_range)) { in __anon5f0f96000502()
629 0, c->UnknownShapeOfRank(rank - dimensions_to_reduce.size())); in __anon5f0f96000502()
663 std::vector<int64_t> dimensions_to_reduce; in __anon5f0f96000702() local
665 c->GetAttr("dimensions_to_reduce", &dimensions_to_reduce)); in __anon5f0f96000702()
[all …]
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dreduction_ops.cc133 const std::vector<int64_t>& dimensions_to_reduce) override { in BuildFinalizer() argument
134 if (dimensions_to_reduce.empty()) { in BuildFinalizer()
137 auto divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]); in BuildFinalizer()
138 for (int i = 1; i < dimensions_to_reduce.size(); i++) { in BuildFinalizer()
139 auto size = xla::GetDimensionSize(input, dimensions_to_reduce[i]); in BuildFinalizer()
Dquantize_and_dequantize_op.cc79 std::vector<int64_t> dimensions_to_reduce; in Compile() local
88 dimensions_to_reduce.reserve(input_rank - 1); in Compile()
91 dimensions_to_reduce.push_back(i); in Compile()
95 dimensions_to_reduce); in Compile()
97 dimensions_to_reduce); in Compile()
Dreduction_ops.h58 const std::vector<int64_t>& dimensions_to_reduce);
/external/tensorflow/tensorflow/compiler/tests/
Dxla_ops_test.py350 x, init_value=0, dimensions_to_reduce=dims, reducer=sum_reducer)
379 x, init_value=1, dimensions_to_reduce=dims, reducer=mul_reducer)
414 dimensions_to_reduce=dims,
419 dimensions_to_reduce=dims,
489 def reduce(values, *, dimensions_to_reduce): argument
494 dimensions_to_reduce=dimensions_to_reduce,
500 dimensions_to_reduce=dimensions_to_reduce,
505 functools.partial(reduce, dimensions_to_reduce=(0,)),
511 functools.partial(reduce, dimensions_to_reduce=(1,)),
517 functools.partial(reduce, dimensions_to_reduce=(0, 1)),
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emission_utils.cc259 llvm::SmallVector<int64_t> dimensions_to_reduce; in IsReductionFromOrToContiguousDimensions() local
261 dimensions_to_reduce.push_back(d.getZExtValue()); in IsReductionFromOrToContiguousDimensions()
265 dimensions_to_reduce); in IsReductionFromOrToContiguousDimensions()
303 llvm::SmallVector<int64_t> dimensions_to_reduce; in GetReductionKindAndContiguousComponents() local
306 dimensions_to_reduce.push_back(d.getZExtValue()); in GetReductionKindAndContiguousComponents()
309 dimensions_to_reduce); in GetReductionKindAndContiguousComponents()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h156 absl::Span<const int64_t> dimensions_to_reduce) override;
Dmlir_hlo_builder.cc227 absl::Span<const int64_t> dimensions_to_reduce) { in ReduceInternal() argument
234 GetI64ElementsAttr(dimensions_to_reduce, &builder_)); in ReduceInternal()
/external/tensorflow/tensorflow/compiler/xla/service/
Dshape_inference.h192 absl::Span<const int64_t> dimensions_to_reduce,
Dhlo_instruction.h915 absl::Span<const int64_t> dimensions_to_reduce,
931 absl::Span<const int64_t> dimensions_to_reduce,
940 absl::Span<const int64_t> dimensions_to_reduce,
Dhlo_instruction.cc1563 absl::Span<const int64_t> dimensions_to_reduce, in CreateReduce() argument
1566 shape, {operand, init_value}, dimensions_to_reduce, reduce_computation)); in CreateReduce()
1573 absl::Span<const int64_t> dimensions_to_reduce, in CreateReduce() argument
1580 shape, all_args, dimensions_to_reduce, reduce_computation); in CreateReduce()
1586 absl::Span<const int64_t> dimensions_to_reduce, in CreateReduce() argument
1593 dimensions_to_reduce, reduce_computation); in CreateReduce()
1603 return CreateReduce(shape, inputs, init_values, dimensions_to_reduce, in CreateReduce()
Dshape_inference.cc2218 absl::Span<const int64_t> dimensions_to_reduce, in InferReduceShape() argument
2245 for (int64_t dimension : dimensions_to_reduce) { in InferReduceShape()
2262 for (int64_t dim_to_reduce : dimensions_to_reduce) { in InferReduceShape()
Dhlo_evaluator.cc3712 absl::Span<const int64_t> dimensions_to_reduce(reduce->dimensions()); in HandleReduce() local
3721 operand_shapes, dimensions_to_reduce, in HandleReduce()
3758 for (const int64_t dim : dimensions_to_reduce) { in HandleReduce()
Dhlo_parser.cc2186 optional<std::vector<int64_t>> dimensions_to_reduce; in CreateInstruction() local
2188 &dimensions_to_reduce}; in CreateInstruction()
2205 arg_shapes, *dimensions_to_reduce, in CreateInstruction()
2217 *dimensions_to_reduce, *reduce_computation)); in CreateInstruction()
Dhlo_instructions.cc1149 absl::Span<const int64_t> dimensions_to_reduce, in HloReduceInstruction() argument
1152 dimensions_to_reduce) { in HloReduceInstruction()
Dshape_inference_test.cc67 absl::Span<const int64_t> dimensions_to_reduce) { in ExpectInferredReduceShape() argument
70 {&arg, &f32_}, dimensions_to_reduce, to_apply); in ExpectInferredReduceShape()
Dhlo_instructions.h862 absl::Span<const int64_t> dimensions_to_reduce,
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h699 absl::Span<const int64_t> dimensions_to_reduce);
704 absl::Span<const int64_t> dimensions_to_reduce);
709 absl::Span<const int64_t> dimensions_to_reduce);
1327 absl::Span<const int64_t> dimensions_to_reduce);
1331 absl::Span<const int64_t> dimensions_to_reduce);
2275 absl::Span<const int64_t> dimensions_to_reduce);
2282 absl::Span<const int64_t> dimensions_to_reduce);
Dxla_builder.cc2581 absl::Span<const int64_t> dimensions_to_reduce) { in Reduce() argument
2584 dimensions_to_reduce); in Reduce()
2590 absl::Span<const int64_t> dimensions_to_reduce) { in Reduce() argument
2609 operand_shape_ptrs, dimensions_to_reduce, called_program_shape)); in Reduce()
2611 dimensions_to_reduce); in Reduce()
2618 absl::Span<const int64_t> dimensions_to_reduce) { in ReduceInternal() argument
2623 for (int64_t dim : dimensions_to_reduce) { in ReduceInternal()
4585 absl::Span<const int64_t> dimensions_to_reduce) { in Reduce() argument
4587 dimensions_to_reduce); in Reduce()
4595 absl::Span<const int64_t> dimensions_to_reduce) { in Reduce() argument
[all …]
/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/
Dops.pyi285 dimensions_to_reduce: Sequence[int]) -> XlaOp: ...
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client_test.py1907 dimensions_to_reduce=[0])
1924 dimensions_to_reduce=[dim])
1940 dimensions_to_reduce=dims)
2620 dimensions_to_reduce=[0])
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_ops_n_z.cc3483 op, result_types, inputs, init_values, op.dimensions_to_reduce(), in matchAndRewrite()
3660 op.init_value(), op.dimensions_to_reduce(), op.reducer()); in matchAndRewrite()
3703 if (op.dimensions_to_reduce().size() > input_ty_0.getRank()) { in verify()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dcanonicalize.mlir2016 …// CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) {dimensions_to_reduce = [], operand_segment_sizes…
2017 …%0 = "tf.XlaReduce"(%arg0, %arg1) {dimensions_to_reduce = [], reducer = @sum1} : (tensor<*xbf16>, …
2028 …// CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) {dimensions_to_reduce = [], operand_segment_size…
2029 …%0 = "tf.XlaVariadicReduce"(%arg0, %arg1) {dimensions_to_reduce = [], reducer = @sum2} : (tensor<3…
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/
Dlegalize-tf.mlir6501 …lse, device = "/job:localhost/replica:0/task:0/device:XLA_GPU:0", dimensions_to_reduce = [0], oper…
6520 …lse, device = "/job:localhost/replica:0/task:0/device:XLA_GPU:0", dimensions_to_reduce = [0], oper…
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc7186 GetI64ElementsAttr(op.dimensions_to_reduce())); in matchAndRewrite()

12