Searched defs:dims_to_reduce (Results 1 – 11 of 11) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/benchmarks/ |
| H A D | reduction_benchmark.cc | 44 ArrayRef<int32_t> dims_to_reduce, in GetReductionIR()
64 llvm::ArrayRef<int32_t> dims_to_reduce) { in GetSumF32IR()
71 llvm::ArrayRef<int32_t> dims_to_reduce) { in GetMeanF32IR()
|
| H A D | reduction_benchmark.h | 36 std::array<int32_t, N> dims_to_reduce) { in ReducedTensorShape()
51 std::array<int32_t, N_DIMS_TO_REDUCE> dims_to_reduce) { in GetEigenSumF32Function()
69 std::array<int32_t, N_DIMS_TO_REDUCE> dims_to_reduce) { in GetEigenMeanF32Function()
|
| H A D | fused_reduction_benchmark.cc | 59 std::array<int64_t, 1> dims_to_reduce{0}; in EigenFusedReduction1D() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | ir_emission_utils.cc | 155 const Shape& input_shape, absl::Span<const int64_t> dims_to_reduce) { in GetReductionKindAndContiguousComponentsImpl()
223 const Shape& operand_shape, absl::Span<int64_t const> dims_to_reduce) { in IsReductionFromOrToContiguousDimensionsImpl()
|
| H A D | tree_reduction_rewriter.cc | 187 std::vector<int64_t> dims_to_reduce = {inner_reduced_dimension}; in RewriteReduction() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/vulkan/ops/ |
| H A D | NativeLayerNorm.cpp | 78 std::vector<int64_t> dims_to_reduce; in native_layer_norm() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | segment_reduction_ops_impl.h | 128 Eigen::IndexList<Eigen::type2index<0> > dims_to_reduce; in Compute() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | layer_norm.cpp | 279 std::vector<int64_t> dims_to_reduce; in rms_norm() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/sparse/cuda/ |
| H A D | SparseCsrTensorMath.cu | 791 Tensor _sparse_csr_prod_cuda(const Tensor& input, IntArrayRef dims_to_reduce, bool keepdim, std::op… in _sparse_csr_prod_cuda()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/sparse/ |
| H A D | SparseCsrTensorMath.cpp | 1347 Tensor _sparse_csr_prod_cpu(const Tensor& input, IntArrayRef dims_to_reduce, bool keepdim, std::opt… in _sparse_csr_prod_cpu()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| H A D | lower_tf.cc | 141 llvm::SmallVector<int64_t, 4> dims_to_reduce(rank - 1); in GetBiasAddGradReductionIndices() local
|