| 41 | } |
| 42 | |
| 43 | af_err af_replace(af_array a, const af_array cond, const af_array b) { |
| 44 | try { |
| 45 | const ArrayInfo& ainfo = getInfo(a); |
| 46 | const ArrayInfo& binfo = getInfo(b); |
| 47 | const ArrayInfo& cinfo = getInfo(cond); |
| 48 | |
| 49 | if (cinfo.ndims() == 0) { return AF_SUCCESS; } |
| 50 | |
| 51 | ARG_ASSERT(2, ainfo.getType() == binfo.getType()); |
| 52 | ARG_ASSERT(1, cinfo.getType() == b8); |
| 53 | |
| 54 | DIM_ASSERT(1, ainfo.ndims() >= binfo.ndims()); |
| 55 | DIM_ASSERT(1, cinfo.ndims() == std::min(ainfo.ndims(), binfo.ndims())); |
| 56 | |
| 57 | dim4 adims = ainfo.dims(); |
| 58 | dim4 bdims = binfo.dims(); |
| 59 | dim4 cdims = cinfo.dims(); |
| 60 | |
| 61 | for (int i = 0; i < 4; i++) { |
| 62 | DIM_ASSERT(1, cdims[i] == std::min(adims[i], bdims[i])); |
| 63 | DIM_ASSERT(2, adims[i] == bdims[i] || bdims[i] == 1); |
| 64 | } |
| 65 | |
| 66 | switch (ainfo.getType()) { |
| 67 | case f16: replace<half>(a, cond, b); break; |
| 68 | case f32: replace<float>(a, cond, b); break; |
| 69 | case f64: replace<double>(a, cond, b); break; |
| 70 | case c32: replace<cfloat>(a, cond, b); break; |
| 71 | case c64: replace<cdouble>(a, cond, b); break; |
| 72 | case s32: replace<int>(a, cond, b); break; |
| 73 | case u32: replace<uint>(a, cond, b); break; |
| 74 | case s64: replace<intl>(a, cond, b); break; |
| 75 | case u64: replace<uintl>(a, cond, b); break; |
| 76 | case s16: replace<short>(a, cond, b); break; |
| 77 | case u16: replace<ushort>(a, cond, b); break; |
| 78 | case s8: replace<schar>(a, cond, b); break; |
| 79 | case u8: replace<uchar>(a, cond, b); break; |
| 80 | case b8: replace<char>(a, cond, b); break; |
| 81 | default: TYPE_ERROR(2, ainfo.getType()); |
| 82 | } |
| 83 | } |
| 84 | CATCHALL; |
| 85 | return AF_SUCCESS; |
| 86 | } |
| 87 | |
| 88 | template<typename ArrayType, typename ScalarType> |
| 89 | void replace_scalar(af_array a, const af_array cond, const ScalarType& b) { |