Searched refs:dims_to_reduce (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | tree_reduction_rewriter.cc | 168 std::vector<int64> dims_to_reduce = {inner_reduced_dimension}; in RewriteReduction() local
170 dims_to_reduce.push_back(0); in RewriteReduction()
177 dims_to_reduce, hlo->to_apply())); in RewriteReduction()
|
| D | ir_emission_utils.cc | 217 const Shape& input_shape, absl::Span<const int64> dims_to_reduce) { in GetReductionKindAndContiguousComponentsImpl() argument
220 if (!absl::c_linear_search(dims_to_reduce, dim)) { in GetReductionKindAndContiguousComponentsImpl()
246 PartitionShapeByMiddleDimensions(input_shape, dims_to_reduce); in GetReductionKindAndContiguousComponentsImpl()
|
| /external/tensorflow/tensorflow/python/ops/ragged/ |
| D | ragged_conversion_ops.py | 98 dims_to_reduce = math_ops.range(
101 default_grad = math_ops.reduce_sum(default_grads, axis=dims_to_reduce)
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | segment_reduction_ops_impl.h | 126 Eigen::DSizes<Eigen::DenseIndex, 1> dims_to_reduce; in Compute() local
127 dims_to_reduce[0] = 0; in Compute()
129 Eigen::IndexList<Eigen::type2index<0> > dims_to_reduce; in Compute() local
197 out_slice = in_slice.reduce(dims_to_reduce, Reducer()); in Compute()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | lower_tf.cc | 105 llvm::SmallVector<int64_t, 4> dims_to_reduce(rank - 1); in GetBiasAddGradReductionIndices() local
106 std::iota(dims_to_reduce.begin(), dims_to_reduce.begin() + feature_dim, 0); in GetBiasAddGradReductionIndices()
107 std::iota(dims_to_reduce.begin() + feature_dim, dims_to_reduce.end(), in GetBiasAddGradReductionIndices()
109 return GetI64ElementsAttr(dims_to_reduce, builder); in GetBiasAddGradReductionIndices()
|