| 93 | } |
| 94 | |
| 95 | af_err af_mean_weighted(af_array *out, const af_array in, |
| 96 | const af_array weights, const dim_t dim) { |
| 97 | try { |
| 98 | ARG_ASSERT(3, (dim >= 0 && dim <= 3)); |
| 99 | |
| 100 | af_array output = 0; |
| 101 | const ArrayInfo &iInfo = getInfo(in); |
| 102 | const ArrayInfo &wInfo = getInfo(weights); |
| 103 | af_dtype iType = iInfo.getType(); |
| 104 | af_dtype wType = wInfo.getType(); |
| 105 | |
| 106 | ARG_ASSERT( |
| 107 | 2, |
| 108 | (wType == f32 || |
| 109 | wType == |
| 110 | f64)); /* verify that weights are non-complex real numbers */ |
| 111 | |
| 112 | // FIXME: We should avoid additional copies |
| 113 | af_array w = weights; |
| 114 | if (iInfo.dims() != wInfo.dims()) { |
| 115 | dim4 iDims = iInfo.dims(); |
| 116 | dim4 wDims = wInfo.dims(); |
| 117 | dim4 tDims(1, 1, 1, 1); |
| 118 | for (int i = 0; i < 4; i++) { |
| 119 | ARG_ASSERT(2, wDims[i] == 1 || wDims[i] == iDims[i]); |
| 120 | tDims[i] = iDims[i] / wDims[i]; |
| 121 | } |
| 122 | AF_CHECK( |
| 123 | af_tile(&w, weights, tDims[0], tDims[1], tDims[2], tDims[3])); |
| 124 | } |
| 125 | |
| 126 | switch (iType) { |
| 127 | case f32: |
| 128 | case s32: |
| 129 | case u32: |
| 130 | case s16: |
| 131 | case u16: |
| 132 | case s8: |
| 133 | case u8: |
| 134 | case b8: output = mean<float>(in, w, dim); break; |
| 135 | case f64: |
| 136 | case s64: |
| 137 | case u64: output = mean<double>(in, w, dim); break; |
| 138 | case c32: output = mean<cfloat>(in, w, dim); break; |
| 139 | case c64: output = mean<cdouble>(in, w, dim); break; |
| 140 | case f16: output = mean<half>(in, w, dim); break; |
| 141 | default: TYPE_ERROR(1, iType); |
| 142 | } |
| 143 | |
| 144 | if (w != weights) { AF_CHECK(af_release_array(w)); } |
| 145 | std::swap(*out, output); |
| 146 | } |
| 147 | CATCHALL; |
| 148 | return AF_SUCCESS; |
| 149 | } |
| 150 | |
| 151 | af_err af_mean_all(double *realVal, double *imagVal, const af_array in) { |
| 152 | try { |