From 298ee88bb75ebaab4646293264e661a309803754 Mon Sep 17 00:00:00 2001 From: Junru Shao Date: Fri, 1 Jun 2018 13:51:29 -0700 Subject: [PATCH] Add save/load json in testcases for foreach (#30) --- tests/python/unittest/test_operator.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 10dfabf72d04..8ab6e0ec9865 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5901,6 +5901,10 @@ def verify_foreach(step, in_syms, state_syms, free_syms, out[i] = out[i] * 2 out.extend(states) out = mx.sym.Group(out) + js_1 = out.tojson() + out = mx.sym.load_json(js_1) + js_2 = out.tojson() + assert js_1 == js_2 arr_grads = [] arg_dict = {} arg_grad_dict = {} @@ -6048,6 +6052,10 @@ def step(in1, states): out1[i] = out1[i] out1.extend(out[1]) out = mx.sym.Group(out1) + js_1 = out.tojson() + out = mx.sym.load_json(js_1) + js_2 = out.tojson() + assert js_1 == js_2 data = mx.nd.arange(4).reshape((1, 2, 2)) state = mx.nd.arange(2) @@ -6121,6 +6129,11 @@ def sym_group(out): h2h_barr_grad1 = mx.nd.empty(h2h_barr.shape) out = mx.sym.contrib.foreach(step, data, [init_h, init_c]) out = sym_group(out) + js_1 = out.tojson() + out = mx.sym.load_json(js_1) + js_2 = out.tojson() + assert js_1 == js_2 + e1 = out.bind(ctx=default_context(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, @@ -6152,6 +6165,11 @@ def sym_group(out): unroll_outs.append(mx.sym.expand_dims(h, axis=0)) unroll_outs = mx.sym.concat(*unroll_outs, dim=0) out = mx.sym.Group([unroll_outs, h, c]) + js_1 = out.tojson() + out = mx.sym.load_json(js_1) + js_2 = out.tojson() + assert js_1 == js_2 + e2 = out.bind(ctx=default_context(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'mylstm_i2h_weight': i2h_warr, 'mylstm_h2h_weight': h2h_warr,