| 27 | } |
| 28 | |
| 29 | af_err af_gradient(af_array *grows, af_array *gcols, const af_array in) { |
| 30 | try { |
| 31 | const ArrayInfo &info = getInfo(in); |
| 32 | af_dtype type = info.getType(); |
| 33 | af::dim4 idims = info.dims(); |
| 34 | |
| 35 | DIM_ASSERT(2, info.elements() > 0); |
| 36 | |
| 37 | af_array grad0; |
| 38 | af_array grad1; |
| 39 | AF_CHECK(af_create_handle(&grad0, idims.ndims(), idims.get(), type)); |
| 40 | AF_CHECK(af_create_handle(&grad1, idims.ndims(), idims.get(), type)); |
| 41 | |
| 42 | switch (type) { |
| 43 | case f32: gradient<float>(&grad0, &grad1, in); break; |
| 44 | case c32: gradient<cfloat>(&grad0, &grad1, in); break; |
| 45 | case f64: gradient<double>(&grad0, &grad1, in); break; |
| 46 | case c64: gradient<cdouble>(&grad0, &grad1, in); break; |
| 47 | default: TYPE_ERROR(1, type); |
| 48 | } |
| 49 | std::swap(*grows, grad0); |
| 50 | std::swap(*gcols, grad1); |
| 51 | } |
| 52 | CATCHALL; |
| 53 | |
| 54 | return AF_SUCCESS; |
| 55 | } |