| 215 | } |
| 216 | |
| 217 | void test_simple_grad(bool grad_cond_out) { |
| 218 | auto graph = ComputingGraph::make(); |
| 219 | HostTensorGenerator<> gen; |
| 220 | auto host_x = gen({2, 3}), host_y = gen({2, 3}), host_pred = gen({1}); |
| 221 | |
| 222 | host_pred->ptr<float>()[0] = 0; |
| 223 | |
| 224 | auto x = opr::Host2DeviceCopy::make(*graph, host_x).rename("x"), |
| 225 | y = opr::Host2DeviceCopy::make(*graph, host_y).rename("y"), |
| 226 | pred = opr::Host2DeviceCopy::make(*graph, host_pred); |
| 227 | auto branches = opr::CondExecPred::make( |
| 228 | pred, |
| 229 | {pred.make_scalar(0.f), pred.make_scalar(1.f), pred.make_scalar(2.f)}); |
| 230 | using GradMode = opr::CondExecMark::Param::GradMode; |
| 231 | auto get_marked = [&branches, grad_cond_out](SymbolVar x, size_t br) { |
| 232 | SymbolVar ret; |
| 233 | unpack_vector( |
| 234 | opr::CondExecMark::make( |
| 235 | branches.at(br), {x}, |
| 236 | {grad_cond_out ? GradMode::SUM_COND_OUT : GradMode::SUM}), |
| 237 | ret); |
| 238 | return ret; |
| 239 | }; |
| 240 | int call_x = 0, call_y = 0; |
| 241 | auto cond_x0 = get_marked(x, 0).rename("cx0"), |
| 242 | cond_x1 = get_marked(x, 1).rename("cx1"), |
| 243 | cond_y = get_marked(y, 2).rename("cy"), |
| 244 | z = merge_one_out( |
| 245 | {cond_x0 * 2, cond_x1 * 3, cond_y * 2.3f}, |
| 246 | MergeMode::EXACT_ONE_SAME_SHAPE) |
| 247 | .rename("merged"), |
| 248 | loss = opr::reduce_sum_sqr(z + y, z.make_scalar(1)), |
| 249 | gx = make_call_rec(cg::grad(loss, x), &call_x), |
| 250 | gy = make_call_rec(cg::grad(loss, y), &call_y); |
| 251 | |
| 252 | std::array<float, 3> kx_all{2.f, 3.f, 0.f}, ky_all{1.f, 1.f, 3.3f}; |
| 253 | |
| 254 | auto make_expect = [&](float kx, float ky, int wrt) { |
| 255 | HostTensorND ret{host_x->comp_node(), host_x->shape()}; |
| 256 | auto pr = ret.ptr<float>(), px = host_x->ptr<float>(), |
| 257 | py = host_y->ptr<float>(); |
| 258 | for (size_t i = 0, it = ret.shape().total_nr_elems(); i < it; ++i) { |
| 259 | float s = px[i] * kx + py[i] * ky, ls = 2 * s; |
| 260 | pr[i] = ls * (wrt ? ky : kx); |
| 261 | } |
| 262 | return ret; |
| 263 | }; |
| 264 | |
| 265 | HostTensorND host_gx, host_gy; |
| 266 | auto func = graph->compile( |
| 267 | {make_callback_copy(gx, host_gx), make_callback_copy(gy, host_gy)}); |
| 268 | |
| 269 | for (size_t i = 0; i < 6; ++i) { |
| 270 | *host_x = *gen({i + 3, 3}); |
| 271 | *host_y = *gen({i + 3, 3}); |
| 272 | |
| 273 | int br_num = i % 3; |
| 274 | host_pred->ptr<float>()[0] = br_num; |
no test coverage detected