| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import pytest |
| import mxnet as mx |
| from numpy.testing import assert_allclose |
| from mxnet.test_utils import * |
| from mxnet.base import _as_list |
| |
| mx.npx.reset_np() |
| |
| def _verify_while_loop(cond, func, loop_var_shapes, free_var_shapes, is_train, max_iterations, is_for, n_steps): |
| |
| def _create_vars(num, prefix): |
| return [mx.sym.var(prefix + str(i)) for i in range(num)] |
| |
| def _create_arrays(shapes): |
| return [mx.nd.random.uniform(-1.0, 1.0, shape=x) for x in shapes] |
| |
| def _create_dict(prefix, shapes): |
| return {prefix + str(i): mx.nd.random.uniform(-1.0, 1.0, shape=x) for i, x in enumerate(shapes)} |
| |
| def _merge_dict(*dicts): |
| result = {} |
| for item in dicts: |
| result.update(item) |
| return result |
| |
| def _to_numpy_list(arrays): |
| return [x.asnumpy() if x is not None else x for x in arrays] |
| |
| def _get_imperative_result(n_steps): |
| free_vars = [args["FreeVar" + str(i)].copy() for i, _ in enumerate(free_var_shapes)] |
| loop_vars = [args["LoopVar" + str(i)].copy() for i, _ in enumerate(loop_var_shapes)] |
| loop_var_start = int(is_for) |
| if is_train: |
| for var in free_vars + loop_vars[loop_var_start: ]: |
| var.attach_grad() |
| with mx.autograd.record(train_mode=is_train): |
| outputs, final_loop_vars = mx.nd.contrib.while_loop( |
| cond=lambda *_loop_vars: cond(_loop_vars, free_vars), |
| func=lambda *_loop_vars: func(_loop_vars, free_vars), |
| loop_vars=loop_vars, |
| max_iterations=max_iterations, |
| ) |
| outputs = _as_list(outputs) |
| final_loop_vars = _as_list(final_loop_vars) |
| outputs = [x[: n_steps] for x in outputs] |
| out_grads = _create_arrays(x.shape for x in outputs) \ |
| + _create_arrays(x.shape for x in final_loop_vars) |
| loop_result_nd = [x * 2 for x in outputs] + [x * 3 for x in final_loop_vars] |
| grads = [] |
| if is_train: |
| cat_out = mx.nd.concat(*[x.reshape(-1) for x in loop_result_nd], dim=0) |
| cat_out.backward(out_grad=mx.nd.concat(*[x.reshape(-1) for x in out_grads], dim=0)) |
| grads = [free_vars[i].grad for i, _ in enumerate(free_var_shapes)] \ |
| + [loop_vars[i].grad for i, _ in enumerate(loop_var_shapes) if i >= loop_var_start] |
| return _to_numpy_list(loop_result_nd), _to_numpy_list(grads), out_grads |
| |
| def _get_symbolic_result(out_grads, n_steps): |
| |
| def _copy_args_dict(name_list): |
| return {name: args[name].copy() for name in name_list} |
| |
| def _zeros_like_dict(name_list): |
| return {name: mx.nd.zeros_like(args[name]) for name in name_list} |
| |
| free_syms = _create_vars(len(free_var_shapes), "FreeVar") |
| loop_syms = _create_vars(len(loop_var_shapes), "LoopVar") |
| outputs, final_loop_syms = mx.sym.contrib.while_loop( |
| cond=lambda *_loop_vars: cond(_loop_vars, free_syms), |
| func=lambda *_loop_vars: func(_loop_vars, free_syms), |
| loop_vars=loop_syms, |
| max_iterations=max_iterations, |
| ) |
| outputs = _as_list(outputs) |
| final_loop_syms = _as_list(final_loop_syms) |
| if n_steps == 0: |
| outputs = [] |
| else: |
| outputs = [x.slice_axis(axis=0, begin=0, end=n_steps) for x in outputs] |
| loop_result_sym = [x * 2 for x in outputs] + [x * 3 for x in final_loop_syms] |
| loop_result_sym = mx.sym.Group(loop_result_sym) |
| |
| loop_var_start = int(is_for) |
| args_names = ["FreeVar" + str(i) for i, _ in enumerate(free_var_shapes)] \ |
| + ["LoopVar" + str(i) for i, _ in enumerate(loop_var_shapes) if i >= loop_var_start] |
| args_grad = None if not is_train else _zeros_like_dict(x for x in args_names) |
| executor = loop_result_sym._bind( |
| ctx=default_device(), |
| args=_copy_args_dict(loop_result_sym.list_inputs()), |
| args_grad=args_grad, |
| ) |
| loop_result_nd = executor.forward(is_train=is_train) |
| grads = [] |
| if is_train: |
| executor.backward(out_grads=out_grads) |
| grads = [executor.grad_dict.get("FreeVar" + str(i), None) for i, _ in enumerate(free_var_shapes)] \ |
| + [executor.grad_dict.get("LoopVar" + str(i), None) for i, _ in enumerate(loop_var_shapes) if i >= loop_var_start] |
| return _to_numpy_list(loop_result_nd), _to_numpy_list(grads) |
| |
| args = _merge_dict( |
| _create_dict("FreeVar", free_var_shapes), |
| _create_dict("LoopVar", loop_var_shapes), |
| ) |
| if is_for: |
| assert loop_var_shapes[0] == (1, ) |
| args["LoopVar0"] = mx.nd.array([0]) |
| imp_outs, imp_grads, out_grads = _get_imperative_result(n_steps) |
| sym_outs, sym_grads = _get_symbolic_result(out_grads, n_steps) |
| for imp_out, sym_out in zip(imp_outs, sym_outs): |
| if imp_out is None or sym_out is None: |
| continue |
| assert_almost_equal(imp_out, sym_out, rtol=1e-3, atol=1e-3) |
| for imp_grad, sym_grad in zip(imp_grads, sym_grads): |
| if imp_grad is None or sym_grad is None: |
| continue |
| assert_almost_equal(imp_grad, sym_grad, rtol=1e-3, atol=1e-3) |
| |
| |
| @pytest.mark.skip(reason="Bug in while loop op, tracked at incubator-mxnet/issues/18575") |
| def test_while_loop_for_foreach(): |
| |
| def make_true_cond(): |
| return lambda loop_vars, _: (loop_vars[0] < 1e35).prod() |
| |
| def make_false_cond(): |
| return lambda loop_vars, _: (loop_vars[0] > 1e35).prod() |
| |
| def make_for_cond(length): |
| return lambda loop_vars, _: loop_vars[0] < length |
| |
| def case_0(): |
| # This is a simple testcase that all loop steps are independent' |
| # It basically scans the array and outputs itself |
| # There is 1 output |
| # There is 1 state: i |
| def _simple_func(loop, free): |
| (i, ), (scanned, ) = loop, free |
| in_ = scanned.take(i).squeeze(axis=0) |
| return (in_, i + 1) |
| _verify_while_loop( |
| cond=make_true_cond(), |
| func=_simple_func, |
| max_iterations=1, |
| is_train=True, |
| is_for=True, |
| loop_var_shapes=[ |
| (1, ), # i |
| ], |
| free_var_shapes=[ |
| (1, 3), # scanned |
| ], |
| n_steps=1, |
| ) |
| |
| def case_1(**params): |
| # This is a simple testcase that simulates a cumulative sum |
| # There is 1 output |
| # There is 1 state: s |
| step_funcs = [ |
| lambda a, b, s: s, |
| lambda a, b, s: a * 1.5 + b * 2.5 - s * 3.5, |
| lambda a, b, s: a * 1.5 - s * 3.5 + b * 2.5, |
| lambda a, b, s: b * 2.5 + a * 1.5 - s * 3.5, |
| lambda a, b, s: b * 2.5 - s * 3.5 + a * 1.5, |
| lambda a, b, s: s * -3.5 + a * 1.5 + b * 2.5, |
| lambda a, b, s: s * -3.5 + b * 2.5 + a * 1.5, |
| lambda a, b, s: a * 2.5 * b + s * 0.3, |
| lambda a, b, s: b * 2.5 * a + s * 0.3, |
| lambda a, b, s: 2.5 * a * b + s * 0.3, |
| lambda a, b, s: b * a * 2.5 + s * 0.3, |
| lambda a, b, s: 2.5 * b * a + s * 0.3, |
| lambda a, b, s: b * a * 2.5 + s * 0.3, |
| lambda a, b, s: s * 0.3 + a * 2.5 * b, |
| lambda a, b, s: s * 0.3 + b * 2.5 * a, |
| lambda a, b, s: s * 0.3 + 2.5 * a * b, |
| lambda a, b, s: s * 0.3 + b * a * 2.5, |
| lambda a, b, s: s * 0.3 + 2.5 * b * a, |
| lambda a, b, s: s * 0.3 + b * a * 2.5, |
| ] |
| def make_func(step_func): |
| def step(loop, free): |
| (s, ), (a, b) = loop, free |
| out = step_func(a, b, s) |
| return (out, out) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| is_train=is_train, |
| is_for=False, |
| **params |
| ) |
| |
| def case_2(**params): |
| # This is a testcase that involves non-differentiable operators |
| # There is 1 output |
| # There is 2 states: i, s |
| step_funcs = [ |
| lambda in_, s, f_1: (in_ * 2) * s * f_1, |
| lambda in_, s, f_1: (in_ * 2) * f_1 * s, |
| lambda in_, s, f_1: s * (in_ * 2) * f_1, |
| lambda in_, s, f_1: s * f_1 * (in_ * 2), |
| lambda in_, s, f_1: f_1 * (in_ * 2) * s, |
| lambda in_, s, f_1: f_1 * s * (in_ * 2), |
| lambda in_, s, f_1: (2 * in_) * s * f_1, |
| lambda in_, s, f_1: (2 * in_) * f_1 * s, |
| lambda in_, s, f_1: s * (2 * in_) * f_1, |
| lambda in_, s, f_1: s * f_1 * (2 * in_), |
| lambda in_, s, f_1: f_1 * (2 * in_) * s, |
| lambda in_, s, f_1: f_1 * s * (2 * in_), |
| ] |
| def make_func(step_func): |
| """This simulates: |
| def compute(s, inputs, f_1, length): |
| outputs = [] |
| for i in range(length): |
| s += inputs[i] * 2 + f_1 |
| outputs.append(s) |
| return outputs, s |
| """ |
| def step(loop, free): |
| (i, s), (scanned, f_1, _) = loop, free |
| in_ = scanned.take(i).squeeze(axis=0) |
| out = step_func(in_, s, f_1) |
| return (out, (i + 1, out)) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| max_iterations=1000, |
| is_train=is_train, |
| is_for=True, |
| **params |
| ) |
| |
| def case_3(length, **params): |
| # This is a testcase for multiple non-differentiable operators and different ways of slicing |
| # There are 2 outputs |
| # There are 3 states: i, s_0, s_1 |
| step_funcs = [ |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_1, |
| lambda i_0, i_1, s_0, s_1, f_0: s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: s_1, |
| lambda i_0, i_1, s_0, s_1, f_0: f_0, |
| ] |
| def make_func(step_func): |
| """This simulates: |
| def compute(input_0, input_1, s_0, s_1, f_0, length): |
| output_0 = [] |
| output_1 = [] |
| for i in range(length): |
| i_0 = input_0[i] |
| i_1 = input_1[length - 1 - i] |
| out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0 |
| s_0 = (s_0 + out) * 1.05 |
| s_1 = (s_1 - out * 0.5) * 0.95 |
| output_0.append(out) |
| output_1.append(out * 1.5) |
| return outputs, s_0, s_1 |
| """ |
| def step(loop, free): |
| (i, s_0, s_1), (sc_0, sc_1, f_0, _) = loop, free |
| i_0 = sc_0.take(i).squeeze(axis=0) |
| i_1 = sc_1.take(length - 1 - i).squeeze(axis=0) |
| out = step_func(i_0, i_1, s_0, s_1, f_0) |
| return ([out, out * 1.5], [i + 1, (s_0 + out) * 1.05, (s_1 - out * 0.5) * 0.95]) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| max_iterations=1000, |
| is_train=is_train, |
| is_for=True, |
| **params |
| ) |
| |
| def case_4(length, single_shape, **params): |
| # It is for the case that inputs & outputs are the same |
| # There are 3 outputs |
| # There are 4 states: i, s_0, s_1, s_2 |
| # i is used in both non-differentiable (take) and differentiable (+) occasions |
| step_funcs = [ |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_1, |
| lambda i_0, i_1, s_0, s_1, f_0: s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: s_1, |
| lambda i_0, i_1, s_0, s_1, f_0: f_0, |
| ] |
| def make_func(step_func): |
| """This simulates: |
| def compute(input_0, input_1, s_0, s_1, s_2, f_0, length): |
| # here s_2 remains untouched |
| output_0 = [] |
| output_1 = [] |
| output_2 = [] |
| for i in range(length): |
| i_0 = input_0[i] |
| i_1 = input_1[length - 1 - i] |
| out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0 |
| out = out * i * i_0 * i_1 |
| s_0 = (s_0 + out) * 1.05 |
| s_1 = (s_1 - out * 0.5) * 0.95 |
| output_0.append(out) |
| output_1.append(f_0) |
| output_2.append(out * 1.5) |
| return output_0, output_1, output_2, s_0, s_1, s_2 |
| """ |
| def step(loop, free): |
| (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free |
| i_0 = sc_0.take(i).squeeze(axis=0) |
| i_1 = sc_1.take(length - 1 - i).squeeze(axis=0) |
| out = step_func(i_0, i_1, s_0, s_1, f_0) |
| out = out * i.reshape([1] * len(single_shape)).broadcast_to(single_shape) |
| out = out * i_0 * i_1 |
| return ([out, f_0, out * 1.5], [i + 1, (s_0 + out) * 1.05, (s_1 - out * 0.5) * 0.95, s_2]) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| max_iterations=1000, |
| is_train=is_train, |
| is_for=True, |
| **params |
| ) |
| |
| def case_5(length, single_shape, **params): |
| # It is for the case that inputs & outputs are the same |
| # There are 0 outputs |
| # There are 4 states: i, s_0, s_1, s_2 |
| # i is used in both differentiable (take) and non-differentiable (+) occasions |
| step_funcs = [ |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_1, |
| lambda i_0, i_1, s_0, s_1, f_0: s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: s_1, |
| lambda i_0, i_1, s_0, s_1, f_0: f_0, |
| ] |
| def make_func(step_func): |
| """This simulates: |
| def compute(input_0, input_1, s_0, s_1, s_2, f_0, length): |
| # here s_2 remains untouched |
| output_0 = [] |
| output_1 = [] |
| output_2 = [] |
| for i in range(length): |
| i_0 = input_0[i] |
| i_1 = input_1[length - 1 - i] |
| out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0 |
| out = out * i * i_0 * i_1 |
| s_0 = (s_0 + out) * 1.05 |
| s_1 = (s_1 - out * 0.5) * 0.95 |
| output_0.append(out) |
| output_1.append(f_0) |
| output_2.append(out * 1.5) |
| return output_0, output_1, output_2, s_0, s_1, s_2 |
| """ |
| def step(loop, free): |
| (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free |
| i_0 = sc_0.take(i).squeeze(axis=0) |
| i_1 = sc_1.take(length - 1 - i).squeeze(axis=0) |
| out = step_func(i_0, i_1, s_0, s_1, f_0) |
| out = out * i.reshape([1] * len(single_shape)).broadcast_to(single_shape) |
| out = out * i_0 * i_1 |
| return ([], [i + 1, (s_0 + out) * 1.05, (s_1 - out * 0.5) * 0.95, s_2]) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| max_iterations=1000, |
| is_train=is_train, |
| is_for=True, |
| **params |
| ) |
| |
| def case_6(length, single_shape, **params): |
| # It is for the case that inputs & outputs are the same |
| # There are 3 outputs |
| # There are 4 states: i, s_0, s_1, s_2 |
| # i is used in both differentiable (take) and non-differentiable (+) occasions |
| step_funcs = [ |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0 * (i_1 * 2) * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * (s_1 * 2) * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * s_0 * f_0 * (s_1 * 2), |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * s_0 * f_0, |
| lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2) * i_0 * (s_1 * 2) * f_0 * s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_0, |
| lambda i_0, i_1, s_0, s_1, f_0: i_1, |
| lambda i_0, i_1, s_0, s_1, f_0: s_0, |
| lambda i_0, i_1, s_0, s_1, f_0: s_1, |
| lambda i_0, i_1, s_0, s_1, f_0: f_0, |
| ] |
| def make_func(step_func): |
| """This simulates: |
| def compute(input_0, input_1, s_0, s_1, s_2, f_0, length): |
| # here s_2 remains untouched |
| output_0 = [] |
| output_1 = [] |
| output_2 = [] |
| for i in range(length): |
| i_0 = input_0[i] |
| i_1 = input_1[length - 1 - i] |
| out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0 |
| out = out * i * i_0 * i_1 |
| s_0 = (s_0 + out) * 1.05 |
| s_1 = (s_1 - out * 0.5) * 0.95 |
| output_0.append(out) |
| output_1.append(f_0) |
| output_2.append(out * 1.5) |
| return output_0, output_1, output_2, s_0, s_1, s_2 |
| """ |
| def step(loop, free): |
| (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free |
| F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd |
| i_0 = sc_0.take(i).squeeze(axis=0) |
| i_1 = sc_1.take(length - 1 - i).squeeze(axis=0) |
| out_0 = step_func(i_0, i_1, s_0, s_1, f_0) |
| out_0 = out_0 * i.reshape([1] * len(single_shape)).broadcast_to(single_shape) |
| out_1 = step_func(i_1, s_0, f_0, s_1, i_0) |
| out_1 = out_1 * i.reshape([1] * len(single_shape)).broadcast_to(single_shape) |
| return ([F.dot(out_0, s_2), f_0, F.dot(s_2, out_1) * 1.5], [i + 1, (s_0 + out_1) * 1.05, (s_1 - out_0 * 0.5) * 0.95, s_2]) |
| return step |
| case_id = 0 |
| for is_train in [True, False]: |
| for step_func in step_funcs: |
| case_id += 1 |
| _verify_while_loop( |
| func=make_func(step_func), |
| max_iterations=1000, |
| is_train=is_train, |
| is_for=True, |
| **params |
| ) |
| |
| # Case 0: the simpest case |
| case_0() |
| # Case 1.1.* |
| case_1( |
| cond=make_true_cond(), |
| loop_var_shapes=[ |
| (1, ), # s |
| ], |
| free_var_shapes=[ |
| (1, ), # a |
| (1, ), # b |
| ], |
| max_iterations=5, |
| n_steps=5, |
| ) |
| # Case 1.2.* |
| case_1( |
| cond=make_true_cond(), |
| loop_var_shapes=[ |
| (2, 3, 4), # s |
| ], |
| free_var_shapes=[ |
| (2, 3, 4), # a |
| (2, 3, 4), # b |
| ], |
| max_iterations=3, |
| n_steps=3, |
| ) |
| # Case 1.3.* |
| case_1( |
| cond=make_false_cond(), |
| loop_var_shapes=[ |
| (2, 3, 4), # s |
| ], |
| free_var_shapes=[ |
| (2, 3, 4), # a |
| (2, 3, 4), # b |
| ], |
| max_iterations=20, |
| n_steps=0, |
| ) |
| # Case 2.1.* |
| case_2( |
| cond=make_for_cond(length=5), |
| loop_var_shapes=[ |
| (1, ), # i |
| (2, ), # s |
| ], |
| free_var_shapes=[ |
| (100, 2), # scanned |
| (2, ), # f_1 |
| (3, 4, 5, 6), # f_2, unused |
| ], |
| n_steps=5, |
| ) |
| # Case 2.2.* |
| case_2( |
| cond=make_for_cond(length=3), |
| loop_var_shapes=[ |
| (1, ), # i |
| (2, ), # s |
| ], |
| free_var_shapes=[ |
| (30, 2), # scanned |
| (2, ), # f_1 |
| (3, 4, 5, 6), # f_2, unused |
| ], |
| n_steps=3, |
| ) |
| # Case 3.* |
| case_3( |
| length=5, |
| cond=make_for_cond(length=5), |
| loop_var_shapes=[ |
| (1, ), # i |
| (2, ), # s_0 |
| (2, ), # s_1 |
| ], |
| free_var_shapes=[ |
| (30, 2), # sc_0 |
| (30, 2), # sc_1 |
| (2, ), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=5, |
| ) |
| # Case 4.1.* |
| case_4( |
| length=4, |
| cond=make_for_cond(length=4), |
| single_shape=[5], |
| loop_var_shapes=[ |
| (1, ), # i |
| (5, ), # s_0 |
| (5, ), # s_1 |
| (23, 6, 8), # s_2 |
| ], |
| free_var_shapes=[ |
| (30, 5), # sc_0 |
| (30, 5), # sc_1 |
| (5, ), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=4, |
| ) |
| # Case 4.2.* |
| case_4( |
| length=5, |
| cond=make_for_cond(length=5), |
| single_shape=[5, 12], |
| loop_var_shapes=[ |
| (1, ), # i |
| (5, 12), # s_0 |
| (5, 12), # s_1 |
| (23, 6, 8), # s_2 |
| ], |
| free_var_shapes=[ |
| (30, 5, 12), # sc_0 |
| (30, 5, 12), # sc_1 |
| (5, 12), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=5, |
| ) |
| # Case 5.1.* |
| case_5( |
| length=4, |
| cond=make_for_cond(length=4), |
| single_shape=[5], |
| loop_var_shapes=[ |
| (1, ), # i |
| (5, ), # s_0 |
| (5, ), # s_1 |
| (23, 6, 8), # s_2 |
| ], |
| free_var_shapes=[ |
| (30, 5), # sc_0 |
| (30, 5), # sc_1 |
| (5, ), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=4, |
| ) |
| # Case 5.2.* |
| case_5( |
| length=5, |
| cond=make_for_cond(length=5), |
| single_shape=[3, 4, 2], |
| loop_var_shapes=[ |
| (1, ), # i |
| (3, 4, 2), # s_0 |
| (3, 4, 2), # s_1 |
| (23, 6, 8), # s_2 |
| ], |
| free_var_shapes=[ |
| (30, 3, 4, 2), # sc_0 |
| (30, 3, 4, 2), # sc_1 |
| (3, 4, 2), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=5, |
| ) |
| # Case 6.* |
| case_6( |
| length=5, |
| cond=make_for_cond(length=5), |
| single_shape=[5, 3], |
| loop_var_shapes=[ |
| (1, ), # i |
| (5, 3), # s_0 |
| (5, 3), # s_1 |
| (3, 5), # s_2 |
| ], |
| free_var_shapes=[ |
| (30, 5, 3), # sc_0 |
| (30, 5, 3), # sc_1 |
| (5, 3), # f_0 |
| (3, 4, 5, 6), # f_1, unused |
| ], |
| n_steps=5, |
| ) |
| |
| |
| def test_while_loop_nested(): |
| |
| def _to_np_list(arrays): |
| return [x.asnumpy() if x is not None else x for x in arrays] |
| |
| def _array(shape): |
| return mx.nd.random.uniform(-1.0, 1.0, shape=shape) |
| |
| def inner_cond(i, j, x_sum, sc): |
| return j < 2 |
| |
| def inner_body(i, j, x_sum, sc): |
| x_ij = sc.take(j).squeeze(axis=0) |
| return (x_ij, x_ij), (i, j + 1, x_sum, sc) |
| |
| def outer_cond(i, j, x_sum, sc): |
| return i < 2 |
| |
| def outer_body(i, j, x_sum, sc): |
| F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd |
| (x_ij, x_ji), (i_p, j_p, x_sum_p, sc_p) = F.contrib.while_loop( |
| cond=inner_cond, |
| func=inner_body, |
| loop_vars=(i, j, x_sum, sc), |
| max_iterations=2, |
| ) |
| return (x_ij, x_ji), (i_p + 1, j_p - 2, x_sum_p, sc_p) |
| |
| def make_loop(i, j, x_sum, sc): |
| F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd |
| (x_ij, x_ji), (new_i, new_j, x_sum_p, sc_p) = F.contrib.while_loop( |
| cond=outer_cond, |
| func=outer_body, |
| loop_vars=(i, j, x_sum, sc), |
| max_iterations=2, |
| ) |
| return new_i, new_j, x_sum_p, sc_p, x_ij, x_ji |
| |
| args = { |
| "i": mx.nd.array([0]), |
| "j": mx.nd.array([0]), |
| "x_sum": _array([5, 3]), |
| "sc": _array([10, 10, 5, 3]), |
| } |
| args_grad = { |
| "x_sum": _array([5, 3]), |
| "sc": _array([10, 10, 5, 3]), |
| } |
| out_grad = [ |
| _array([1]), |
| _array([1]), |
| _array([5, 3]), |
| _array([10, 10, 5, 3]), |
| _array([2, 2, 10, 5, 3]), |
| _array([2, 2, 10, 5, 3]), |
| ] |
| def _get_imp_result(is_train, args, args_grad, out_grad): |
| args = {k: v.copy() for k, v in args.items()} |
| args_grad = {k: v.copy() for k, v in args_grad.items()} |
| i, j, x_sum, sc = [args[x].copy() for x in ["i", "j", "x_sum", "sc"]] |
| if is_train: |
| x_sum.attach_grad() |
| sc.attach_grad() |
| with mx.autograd.record(train_mode=is_train): |
| results = make_loop(i, j, x_sum, sc) |
| cat_res = mx.nd.concat(*[x.reshape(-1) for x in results], dim=0) |
| if not is_train: |
| return _to_np_list(results), [] |
| cat_grad = mx.nd.concat(*[x.reshape(-1) for x in out_grad], dim=0) |
| assert cat_grad.shape == cat_res.shape |
| cat_res.backward(out_grad=cat_grad) |
| grads = [x_sum.grad, sc.grad] |
| return _to_np_list(results), _to_np_list(grads) |
| |
| def _get_sym_result(is_train, args, args_grad, out_grad): |
| args = {k: v.copy() for k, v in args.items()} |
| args_grad = {k: v.copy() for k, v in args_grad.items()} |
| i, j, x_sum, sc = [ |
| mx.sym.var("i"), |
| mx.sym.var("j"), |
| mx.sym.var("x_sum"), |
| mx.sym.var("sc"), |
| ] |
| result_sym = mx.sym.Group(make_loop(i, j, x_sum, sc)) |
| executor = result_sym._bind( |
| ctx=default_device(), |
| args=args, |
| args_grad=args_grad, |
| ) |
| results = executor.forward(is_train=is_train) |
| if not is_train: |
| return _to_np_list(results), [] |
| executor.backward(out_grads=out_grad) |
| grads = [executor.grad_dict["x_sum"], executor.grad_dict["sc"]] |
| return _to_np_list(results), _to_np_list(grads) |
| |
| for is_train in [True, False]: |
| imp_out, imp_grad = _get_imp_result(is_train=is_train, args=args, args_grad=args_grad, out_grad=out_grad) |
| sym_out, sym_grad = _get_sym_result(is_train=is_train, args=args, args_grad=args_grad, out_grad=out_grad) |
| assert len(imp_out) == len(sym_out) |
| assert len(imp_grad) == len(sym_grad) |
| for x, y in zip(imp_out, sym_out): |
| assert_almost_equal(x, y, rtol=1e-3, atol=1e-3) |
| for x, y in zip(imp_grad, sym_grad): |
| assert_almost_equal(x, y, rtol=1e-3, atol=1e-3) |
| |
| |
| def _verify_cond(cond_func, then_func, else_func, input_var_shapes, free_var_shapes, is_train): |
| |
| def _create_symbol(prefix, i): |
| return mx.sym.var(prefix + str(i)) |
| |
| def _create_array(shape): |
| return mx.nd.random.uniform(-1.0, 1.0, shape=shape) |
| |
| def _to_numpy_list(arrays): |
| return [x.asnumpy() if x is not None else x for x in arrays] |
| |
| def _merge_dict(*dicts): |
| result = {} |
| for item in dicts: |
| result.update(item) |
| return result |
| |
| _input_syms = [_create_symbol("InputVar", i) for i, _ in enumerate(input_var_shapes)] |
| _free_syms = [_create_symbol("FreeVar", i) for i, _ in enumerate(free_var_shapes)] |
| _input_vars = [_create_array(x) for x in input_var_shapes] |
| _free_vars = [_create_array(x) for x in free_var_shapes] |
| _args_dict = _merge_dict( |
| {"InputVar" + str(i): x for i, x in enumerate(_input_vars)}, |
| {"FreeVar" + str(i): x for i, x in enumerate(_free_vars)}, |
| ) |
| |
| def _get_imperative_result(): |
| free_vars = [x.copy() for x in _free_vars] |
| input_vars = [x.copy() for x in _input_vars] |
| out_grads = [] |
| if is_train: |
| for var in free_vars + input_vars: |
| var.attach_grad() |
| with mx.autograd.record(train_mode=is_train): |
| outputs = mx.nd.contrib.cond( |
| pred=cond_func(input_vars, free_vars), |
| then_func=lambda: then_func(input_vars, free_vars), |
| else_func=lambda: else_func(input_vars, free_vars), |
| ) |
| outputs = _as_list(outputs) |
| outputs = [x * 2 for x in outputs] |
| grads = [] |
| if is_train: |
| out_grads = [_create_array(x.shape) for x in outputs] |
| cat_out = mx.nd.concat(*[x.reshape(-1) for x in outputs], dim=0) |
| cat_out.backward(out_grad=mx.nd.concat(*[x.reshape(-1) for x in out_grads], dim=0)) |
| grads = [free_vars[i].grad for i, _ in enumerate(free_var_shapes)] \ |
| + [input_vars[i].grad for i, _ in enumerate(input_var_shapes)] |
| return _to_numpy_list(outputs), _to_numpy_list(grads), out_grads |
| |
| def _get_symbolic_result(out_grads): |
| outputs_sym = mx.sym.contrib.cond( |
| pred=cond_func(_input_syms, _free_syms), |
| then_func=lambda: then_func(_input_syms, _free_syms), |
| else_func=lambda: else_func(_input_syms, _free_syms), |
| ) |
| outputs_sym = _as_list(outputs_sym) |
| outputs_sym = [x * 2 for x in outputs_sym] |
| outputs_sym = mx.sym.Group(outputs_sym) |
| executor = outputs_sym._bind( |
| ctx=default_device(), |
| args={name: _args_dict[name].copy() for name in outputs_sym.list_inputs()}, |
| args_grad=None if not is_train else _merge_dict( |
| {"InputVar" + str(i): mx.nd.zeros(s) for i, s in enumerate(input_var_shapes)}, |
| {"FreeVar" + str(i): mx.nd.zeros(s) for i, s in enumerate(free_var_shapes)}, |
| ), |
| ) |
| outputs = executor.forward(is_train=is_train) |
| grads = [] |
| if is_train: |
| executor.backward(out_grads=out_grads) |
| grads = [executor.grad_dict.get("FreeVar" + str(i), None) for i, _ in enumerate(free_var_shapes)] \ |
| + [executor.grad_dict.get("InputVar" + str(i), None) for i, _ in enumerate(input_var_shapes)] |
| return _to_numpy_list(outputs), _to_numpy_list(grads) |
| |
| imp_outs, imp_grads, out_grads = _get_imperative_result() |
| sym_outs, sym_grads = _get_symbolic_result(out_grads) |
| for imp_out, sym_out in zip(imp_outs, sym_outs): |
| if imp_out is None or sym_out is None: |
| continue |
| assert_almost_equal(imp_out, sym_out, rtol=1e-3, atol=1e-3) |
| for imp_grad, sym_grad in zip(imp_grads, sym_grads): |
| if imp_grad is None or sym_grad is None: |
| continue |
| assert_almost_equal(imp_grad, sym_grad, rtol=1e-3, atol=1e-3) |
| |
| |
| def test_cond(): |
| # whether there are free variables in three graphs |
| # whether these three graphs contain input_vars |
| # whether to use all input_vars |
| # which branch to choose |
| def run_case(cond_func, then_func, else_func, **params): |
| def make_cond(is_inverse): |
| def cond(inputs, free): |
| x = cond_func(inputs, free) |
| if is_inverse: |
| if isinstance(x, mx.sym.Symbol): |
| return mx.sym.logical_not(x) |
| else: |
| return mx.nd.logical_not(x) |
| return x |
| return cond |
| for is_train in [True, False]: |
| for is_inverse in [False, True]: |
| _verify_cond( |
| cond_func=make_cond(is_inverse), |
| then_func=then_func, |
| else_func=else_func, |
| is_train=is_train, |
| **params |
| ) |
| # Each function can |
| # 1. use_free_vars or not: T/F |
| # 2. use_input_vars or not: T/F |
| # 3. use_all_input_vars or not: T/F |
| # (a, b, c) are inputs, (d, e, f) are free_vars |
| cond_funcs = [ |
| lambda a, b, c, d, e, f: (a * b).sum() < 0.5, # F, T, F |
| lambda a, b, c, d, e, f: (a + b + c).sum() < 0.5, # F, T, T |
| lambda a, b, c, d, e, f: (d + e).sum() < 0.5, # T, F, F |
| lambda a, b, c, d, e, f: (d + e * a).sum() < 0.5, # T, T, F |
| lambda a, b, c, d, e, f: (d + e * a + b * c).sum() < 0.5, # T, T, T |
| ] |
| body_funcs = [ |
| lambda a, b, c, d, e, f: a * b, # F, T, F |
| lambda a, b, c, d, e, f: a * b * c, # F, T, T |
| lambda a, b, c, d, e, f: d * e, # T, F, F |
| lambda a, b, c, d, e, f: d * e * a, # T, T, F |
| lambda a, b, c, d, e, f: d * e * a * b * c, # T, T, T |
| # some extra tests |
| lambda a, b, c, d, e, f: b * c, |
| lambda a, b, c, d, e, f: a * c, |
| lambda a, b, c, d, e, f: (a + b) * c, |
| lambda a, b, c, d, e, f: c * (b - a), |
| ] |
| # enumerate all kinds of possible combinations |
| for cond_func in cond_funcs: |
| for then_func in body_funcs: |
| for else_func in body_funcs: |
| run_case( |
| cond_func=lambda x, y: cond_func(x[0], x[1], x[2], y[0], y[1], y[2]), |
| then_func=lambda x, y: then_func(x[0], x[1], x[2], y[0], y[1], y[2]), |
| else_func=lambda x, y: else_func(x[0], x[1], x[2], y[0], y[1], y[2]), |
| input_var_shapes=[ |
| (2, 3), |
| (2, 3), |
| (2, 3), |
| ], |
| free_var_shapes=[ |
| (2, 3), |
| (2, 3), |
| (2, 3), |
| ] |
| ) |
| |
| |
| @pytest.mark.garbage_expected |
| def test_foreach(): |
| v3 = mx.sym.var("v0") |
| v4 = mx.sym.var("v1") |
| v5 = mx.sym.var("v2") |
| v6 = mx.sym.var("v3") |
| v7 = mx.sym.var("v4") |
| v8 = mx.sym.var("v5") |
| |
| def verify_foreach(step, in_syms, state_syms, free_syms, |
| in_arrs, init_states, frees, out_grads, is_train=True, |
| free_vars_func=None, num_iters=1): |
| step_sym = lambda in_syms, state_syms : step(in_syms, state_syms, free_syms) |
| res, states = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) |
| out = _as_list(res) |
| num_outputs = len(out) |
| for i in range(num_outputs): |
| out[i] = out[i] * 2 |
| out.extend(states) |
| out = mx.sym.Group(out) |
| js_1 = out.tojson() |
| out = mx.sym.fromjson(js_1) |
| js_2 = out.tojson() |
| assert js_1 == js_2 |
| arr_grads = [] |
| arg_dict = {} |
| arg_grad_dict = {} |
| i = 0 |
| for arr in _as_list(in_arrs): |
| arr_grad = mx.nd.empty(arr.shape) |
| arr_grads.append(arr_grad) |
| arg_dict['v'+str(i)] = arr |
| arg_grad_dict['v'+str(i)] = arr_grad |
| i = i + 1 |
| for arr in init_states: |
| arr_grad = mx.nd.empty(arr.shape) |
| arr_grads.append(arr_grad) |
| arg_dict['v'+str(i)] = arr |
| arg_grad_dict['v'+str(i)] = arr_grad |
| i = i + 1 |
| for arr in frees: |
| arr_grad = mx.nd.empty(arr.shape) |
| arr_grads.append(arr_grad) |
| arg_dict['v'+str(i)] = arr |
| arg_grad_dict['v'+str(i)] = arr_grad |
| i = i + 1 |
| |
| if is_train: |
| e = out._bind(ctx=default_device(), args=arg_dict, args_grad=arg_grad_dict) |
| else: |
| e = out._bind(ctx=default_device(), args=arg_dict) |
| # the inputs to forward and backward are the same so forward and backward |
| # should always return the same outputs. |
| for _ in range(num_iters): |
| e.forward(is_train=is_train) |
| if (is_train): |
| # backward |
| tmp_grads = out_grads[0][:] |
| tmp_grads.extend(out_grads[1]) |
| e.backward(tmp_grads) |
| |
| # Below we use imperative to reimplement foreach and compute its gradients. |
| res = [] |
| for _ in range(len(_as_list(out_grads[0]))): |
| res.append([]) |
| for arr in _as_list(in_arrs): |
| arr.attach_grad() |
| for arr in init_states: |
| arr.attach_grad() |
| for arr in frees: |
| arr.attach_grad() |
| with mx.autograd.record(): |
| frees_imp = frees if free_vars_func is None else free_vars_func(frees) |
| step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees_imp) |
| states = [mx.nd.expand_dims(s, 0) for s in init_states] |
| res, states = mx.nd.contrib.foreach(step_imp, in_arrs, init_states) |
| |
| res2 = _as_list(res) |
| for i in range(len(res2)): |
| res2[i] = res2[i] * 2 |
| outs = [] |
| outs[:] = res2[:] |
| if isinstance(states, list): |
| outs.extend(states) |
| states = [mx.nd.expand_dims(s, 0) for s in states] |
| res2.extend(states) |
| else: |
| outs.append(states) |
| states = mx.nd.expand_dims(states, 0) |
| res2.append(states) |
| if is_train: |
| res = mx.nd.concat(*res2, dim=0) |
| |
| tmp_grads = out_grads[0][:] |
| tmp_grads1 = [mx.nd.expand_dims(grad, 0) for grad in out_grads[1]] |
| tmp_grads.extend(tmp_grads1) |
| if is_train: |
| res.backward(mx.nd.concat(*tmp_grads, dim=0)) |
| for i in range(len(outs)): |
| assert e.outputs[i].shape == outs[i].shape |
| assert_almost_equal(e.outputs[i].asnumpy(), outs[i].asnumpy(), |
| rtol=1e-3, atol=1e-3) |
| if (is_train): |
| all_ins = _as_list(in_arrs)[:] |
| all_ins.extend(init_states) |
| all_ins.extend(frees) |
| size = min(len(all_ins), len(e.grad_arrays)) |
| for i in range(size): |
| assert_almost_equal(all_ins[i].grad.asnumpy(), |
| e.grad_arrays[i].asnumpy(), |
| rtol=1e-3, atol=1e-3) |
| |
| # Test cases: |
| # * graph inputs are stored in different orders. |
| # This is to test if foreach finds the data arrays and weight arrays |
| # in the right location. |
| # * the number of iterations: odd or even. |
| # * multiple inputs and multiple outputs. |
| # * inference. |
| def step1(in1, states, free): |
| out = in1 * 2 + states[0] + free[0] |
| return (out, [out]) |
| frees1 = [mx.nd.arange(2), mx.nd.arange(2) + 1] |
| arrs = mx.nd.arange(6).reshape(shape=(3, 2)) |
| states = [mx.nd.arange(2)] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, True, |
| lambda frees : [frees[0] + frees[1]]) |
| verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, False, |
| lambda frees : [frees[0] + frees[1]]) |
| verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, True, |
| lambda frees : [frees[0] + frees[1]], 5) |
| verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, False, |
| lambda frees : [frees[0] + frees[1]], 5) |
| |
| # Test the even number of iterations. |
| frees = [mx.nd.random.uniform(shape=(2))] |
| arrs = mx.nd.random.uniform(shape=(2, 2)) |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| # Test the odd number of iterations |
| arrs = mx.nd.random.uniform(shape=(3, 2)) |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| |
| # Reorder the input and state in the subgraph inputs. |
| def step2(in1, states, free): |
| out = states[0] + in1 * 2 + free[0] |
| return (out, [out]) |
| # Test the even number of iterations. |
| arrs = mx.nd.random.uniform(shape=(2, 2)) |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| # Test the odd number of iterations. |
| arrs = mx.nd.random.uniform(shape=(3, 2)) |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| |
| # Test multiple inputs and outputs. |
| def step3(in1, states, free): |
| out = in1[0] + in1[1] * 2 + states[0] + states[1] * 2 + free[0] |
| return ([out, out], [out * 2, out * 3]) |
| arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] |
| states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[1].shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] |
| verify_foreach(step3, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads) |
| verify_foreach(step3, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads, False) |
| |
| # Test multiple inputs and outputs. |
| # The order of subgraph inputs doesn't match the operator inputs |
| def step4(in1, states, free): |
| out = in1[1] * 2 + states[0] + free[0] + states[1] * 2 + in1[0] |
| return ([out, out * 2], [out * 2, out * 3]) |
| arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] |
| states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[1].shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] |
| verify_foreach(step4, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads) |
| verify_foreach(step4, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads, False) |
| |
| # Test multiple inputs and outputs. |
| # The data inputs and states have different shapes. |
| def step5(in1, states, free): |
| if isinstance(in1[0], mx.nd.NDArray): |
| out1 = mx.nd.broadcast_add(states[0] + free[1], in1[1] * 2) |
| out2 = mx.nd.broadcast_add(in1[0], free[0] + states[1] * 2) |
| else: |
| out1 = mx.sym.broadcast_add(states[0] + free[1], in1[1] * 2) |
| out2 = mx.sym.broadcast_add(in1[0], free[0] + states[1] * 2) |
| return ([out1, out2 * 2], [states[0] * 2, states[1] * 3]) |
| frees = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2, 2))] |
| arrs = [mx.nd.random.uniform(shape=(3, 2, 2)), mx.nd.random.uniform(shape=(3, 2))] |
| states = [mx.nd.random.uniform(shape=(2, 2)), mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[0].shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] |
| verify_foreach(step5, [v3, v4], [v5, v6], [v7, v8], arrs, states, frees, out_grads, False) |
| |
| # Test multiple inputs and outputs. |
| # The data inputs and states have different shapes and data types. |
| def step6(in1, states, free): |
| if isinstance(in1[0], mx.nd.NDArray): |
| out1 = mx.nd.broadcast_add(states[0] + mx.nd.cast(free[1], 'float32'), |
| mx.nd.cast(in1[1], 'float32') * 2) |
| out2 = mx.nd.broadcast_add(in1[0], |
| free[0] + mx.nd.cast(states[1], 'float32') * 2) |
| else: |
| out1 = mx.sym.broadcast_add(states[0] + mx.sym.cast(free[1], 'float32'), |
| mx.sym.cast(in1[1], 'float32') * 2) |
| out2 = mx.sym.broadcast_add(in1[0], |
| free[0] + mx.sym.cast(states[1], 'float32') * 2) |
| return ([out1, out2 * 2], [states[0] * 2, states[1] * 3]) |
| frees = [mx.nd.random.uniform(shape=(2)), |
| mx.nd.cast(mx.nd.random.uniform(shape=(2, 2)), 'float64')] |
| arrs = [mx.nd.random.uniform(shape=(3, 2, 2)), |
| mx.nd.cast(mx.nd.random.uniform(shape=(3, 2)), dtype='float16')] |
| states = [mx.nd.random.uniform(shape=(2, 2)), |
| mx.nd.cast(mx.nd.random.uniform(shape=(2)), dtype='int32')] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[0].shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] |
| verify_foreach(step6, [v3, v4], [v5, v6], [v7, v8], arrs, states, frees, out_grads, False) |
| |
| # Test multiple inputs and outputs. |
| # some of the inputs are used twice. |
| def step7(in1, states, free): |
| out1 = states[0] + in1[0] + free[1] + in1[1] * 2 + free[0] |
| out2 = in1[0] + free[0] + states[1] * 2 + in1[1] |
| return ([out1, out2 * 2], [states[0] * 2, states[1] * 3]) |
| frees = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] |
| arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] |
| states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[0].shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] |
| verify_foreach(step7, [v3, v4], [v5, v6], [v7, v8], arrs, states, frees, out_grads, False) |
| |
| # Test the case that the output is the input. |
| arrs = mx.nd.random.uniform(shape=(3, 2)) |
| states = [mx.nd.arange(2)] |
| frees = [mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], |
| [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| def step8(in1, states, free): |
| return (in1, [states[0] * free[0]]) |
| verify_foreach(step8, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step8, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| def step9(in1, states, free): |
| return (in1 * free[0], states) |
| verify_foreach(step9, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step9, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| |
| # Test the case that not all inputs are used. |
| def step10(in1, states, free): |
| return (in1, states) |
| verify_foreach(step10, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step10, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| def step11(in1, states, free): |
| return (in1, free) |
| try: |
| verify_foreach(step11, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step11, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| except AssertionError: |
| print("the states have to be used") |
| def step12(in1, states, free): |
| return (in1, [states[0] + 1, states[0] + 2]) |
| states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] |
| frees = [] |
| try: |
| verify_foreach(step12, v3, [v4, v5], [], arrs, states, frees, out_grads) |
| verify_foreach(step12, v3, [v4, v5], [], arrs, states, frees, out_grads, False) |
| except AssertionError: |
| print("the states have to be used") |
| |
| # test without free variables. |
| def step13(in1, states, free): |
| return (in1, states) |
| states = [mx.nd.random.uniform(shape=(2))] |
| verify_foreach(step13, v3, [v4], [], arrs, states, [], out_grads) |
| verify_foreach(step13, v3, [v4], [], arrs, states, [], out_grads, False) |
| |
| # test when there isn't output data or output states. |
| def step14(in1, states, free): |
| return (in1 + free[0], []) |
| frees = [mx.nd.random.uniform(shape=(2))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], []] |
| verify_foreach(step14, v3, [], [v4], arrs, [], frees, out_grads) |
| verify_foreach(step14, v3, [], [v4], arrs, [], frees, out_grads, False) |
| def step15(in1, states, free): |
| return ([], [in1 * states[0] * free[0]]) |
| out_grads = [[], [mx.nd.random.uniform(-10, 10, states[0].shape)]] |
| verify_foreach(step15, v3, [v4], [v5], arrs, states, frees, out_grads) |
| verify_foreach(step15, v3, [v4], [v5], arrs, states, frees, out_grads, False) |
| |
| # Test the case of iterating on a 1D data array. |
| def step16(in1, states, free): |
| return ([in1[0] * states[0]], [states[0] * 2]) |
| arrs = [mx.nd.arange(3)] |
| states = [mx.nd.random.uniform(shape=(1))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, (3, 1))], |
| [mx.nd.random.uniform(-10, 10, (1))]] |
| verify_foreach(step16, [v3], [v4], [], arrs, states, [], out_grads) |
| verify_foreach(step16, [v3], [v4], [], arrs, states, [], out_grads, False) |
| def step17(in1, states, free): |
| return ([in1[1] * in1[0] * states[0]], [states[0] * 2]) |
| arrs = [mx.nd.random.uniform(shape=(3, 1)), mx.nd.arange(3)] |
| states = [mx.nd.random.uniform(shape=(1))] |
| out_grads = [[mx.nd.random.uniform(-10, 10, (3, 1))], |
| [mx.nd.random.uniform(-10, 10, (1))]] |
| verify_foreach(step17, [v3, v4], [v5], [], arrs, states, [], out_grads) |
| verify_foreach(step17, [v3, v4], [v5], [], arrs, states, [], out_grads, False) |
| |
| |
| def test_foreach_nested(): |
| # Test nested foreach. |
| def step_in(in1, states): |
| out = in1 * 2 + states[0] |
| return (out, [out]) |
| |
| def step_sym(in1, states): |
| out1 = mx.sym.contrib.foreach(step_in, in1, states) |
| out = mx.sym.broadcast_add(out1[0], states[0]) |
| return (out, [mx.sym.squeeze(mx.sym.slice(out, begin=(0, 0), end=(1, 2)))]) |
| def step_nd(in1, states): |
| out1 = mx.nd.contrib.foreach(step_in, in1, states) |
| out = mx.nd.broadcast_add(out1[0], states[0]) |
| return (out, [mx.nd.squeeze(mx.nd.slice(out, begin=(0, 0), end=(1, 2)))]) |
| |
| data_sym = mx.sym.var("v1") |
| state_sym = mx.sym.var("v2") |
| out, states = mx.sym.contrib.foreach(step_sym, data_sym, [state_sym]) |
| assert isinstance(states, list) |
| assert len(states) == 1 |
| out = mx.sym.broadcast_add(out, states[0]) |
| |
| js_1 = out.tojson() |
| out = mx.sym.fromjson(js_1) |
| js_2 = out.tojson() |
| assert js_1 == js_2 |
| |
| data = mx.nd.arange(8).reshape((2, 2, 2)) |
| state = mx.nd.arange(2) |
| data_grad = mx.nd.empty(data.shape) |
| state_grad = mx.nd.empty(state.shape) |
| e = out._bind(ctx=default_device(), args={'v1':data, 'v2':state}, |
| args_grad={'v1':data_grad, 'v2':state_grad}) |
| e.forward(is_train=True) |
| out_grads = [] |
| for out in e.outputs: |
| out_grads.append(mx.nd.random.uniform(shape=out.shape)) |
| e.backward(out_grads) |
| |
| data.attach_grad() |
| state.attach_grad() |
| with mx.autograd.record(): |
| out, states = mx.nd.contrib.foreach(step_nd, data, [state]) |
| assert isinstance(states, list) |
| assert len(states) == 1 |
| res = mx.nd.broadcast_add(out, states[0]) |
| assert_almost_equal(res.asnumpy(), e.outputs[0].asnumpy(), rtol=1e-3, atol=1e-3) |
| |
| res.backward(out_grads[0]) |
| assert_almost_equal(data.grad.asnumpy(), data_grad.asnumpy(), rtol=1e-3, atol=1e-3) |
| assert_almost_equal(state.grad.asnumpy(), state_grad.asnumpy(), rtol=1e-3, atol=1e-3) |
| |
| |
| def test_foreach_with_unkown_dim(): |
| # MXNet supports using 0 as placeholder for unknown dimensions in shape |
| step = lambda data, states: (data + states[0], [states[0] * 2]) |
| # input shape with NCHW format and N is unknown |
| data = mx.sym.var('data', shape=(0, 3, 32, 32)) |
| states = [mx.sym.var('state')] |
| outs, states = mx.sym.contrib.foreach(step, data, states) |
| _, output_shape, _ = outs.infer_shape_partial() |
| assert_allclose((0, 3, 32, 32), output_shape[0]) |
