def test_load_linear_regressor(m_):
shape_dict = {"m": m_}
m = pm.parameter("m")
mu = pm.parameter(name="mu", default=1.0)
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
graph = pm.linear_regressor_train(x, w, y, mu, m)
test_graph, input_info, out_info, keys = linear(m=m_, coarse=True)
assert len(test_graph.nodes.keys()) == len(graph.nodes.keys())
assert op_counts(test_graph) == op_counts(graph)
shape_val_pass = pm.NormalizeGraph(shape_dict)
new_graph = shape_val_pass(graph)
test_res = new_graph(keys, input_info)
np.testing.assert_allclose(test_res, out_info["w"])
test_graph_lowered, input_info, new_out_info, keys = linear(m=m_)
flatten_pass = pm.Lower({})
test_flatten_pass = pm.Lower({})
flattened_g = flatten_pass(new_graph)
ref_lowered = test_flatten_pass(test_graph_lowered, {})
assert len(ref_lowered.nodes.keys()) == len(flattened_g.nodes.keys())
assert op_counts(ref_lowered) == op_counts(flattened_g)
all_vals = flattened_g(keys, input_info)
np.testing.assert_allclose(new_out_info["w"], all_vals)
def test_linear_reg():
m_ = 3
graph, input_info, out_info, keys = linear(m=m_, coarse=True)
coarse_eval = graph(keys, input_info)
np.testing.assert_allclose(coarse_eval, out_info["w"])
fgraph, input_info, out_info, keys = linear(m=m_, coarse=False)
lower_pass = Lower({})
lowered_graph = lower_pass(fgraph, {})
all_vals = lowered_graph(keys, input_info)
out = np.asarray(all_vals).reshape(out_info["w"].shape)
np.testing.assert_allclose(out, out_info["w"])
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
pb_path = f"{full_path}/{graph.name}.srdfg"
pm.pb_store(lowered_graph, full_path)
loaded_node = pm.pb_load(pb_path)
_, input_info, out_info, keys = linear(m=m_, coarse=False)
loaded_res = loaded_node(keys, input_info)
out = np.asarray(loaded_res).reshape(out_info["w"].shape)
np.testing.assert_allclose(out, out_info["w"])
def test_linear(m_):
shape_dict = {"m": m_}
graph, input_info, out_info, keys = linear(**shape_dict, coarse=True)
shape_val_pass = pm.NormalizeGraph(shape_dict)
new_graph = shape_val_pass(graph)
test_res = new_graph(keys, input_info)
np.testing.assert_allclose(test_res, out_info["w"])
graph, input_info, new_out_info, keys = linear(**shape_dict)
flatten_pass = pm.Lower({})
flattened_g = flatten_pass(new_graph)
all_vals = flattened_g(keys, input_info)
np.testing.assert_allclose(new_out_info["w"], all_vals)
def create_linear(m):
shape_dict = {"m": m}
graph, input_info, out_info, keys = linear(m=m, coarse=True)
_, input_info, out_info, keys = linear(m=m, coarse=False)
cwd = Path(f"{__file__}").parent
full_path = f"{cwd}/outputs"
tabla_path = f"{full_path}/{graph.name}_{m}_tabla.json"
tabla_ir, tabla_graph = pm.generate_tabla(graph,
shape_dict,
tabla_path,
context_dict=input_info,
add_kwargs=True)
def test_load_nested_linear_regressor(m_):
shape_dict = {"m": m_}
with pm.Node(name="nested_linear") as graph:
m = pm.parameter(name="m")
mu = pm.parameter(name="mu", default=1.0)
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
pm.linear_regressor_train(x, w, y, mu, m, name="linear_regressor")
j = pm.index(0, m-1, name="j")
tw = (w[j] - 4).set_name("tw")
test_graph, input_info, out_info, keys = linear(m=m_, coarse=True)
shape_val_pass = pm.NormalizeGraph(shape_dict)
new_graph = shape_val_pass(graph)
test_res = new_graph("tw", input_info)
np.testing.assert_allclose(test_res, (out_info["w"] - 4))
ref_graph, input_info, new_out_info, keys = linear(m=m_)
flatten_pass = pm.Lower({})
keys = [f"tw/tw({i},)" for i in range(m_)]
flattened_g = flatten_pass(new_graph)
all_vals = flattened_g(keys, input_info)
def test_translate_linear_regressor(m):
out_key_map = {"y": "y:0", "x": "x:0", "w": "W:0"}
in_key_map = [("w", "W:0")]
fpath = f"{ONNX_FILE_DIR}/linear_{m}.onnx"
shape_dict = {"m": m}
graph = pm.from_onnx(fpath)
test_graph, input_info, out_info, keys = linear(m=m, coarse=True)
tinput_info = copy.deepcopy(input_info)
tkeys = copy.deepcopy(keys)
test_res = test_graph(tkeys, tinput_info)
np.testing.assert_allclose(test_res, (out_info["w"]))
onx_input_info = copy.deepcopy(input_info)
translated_inputs = {out_key_map[k]: v for k,v in input_info.items() if k in out_key_map}
onnx_res = graph(in_key_map[0][1], translated_inputs)
np.testing.assert_allclose(onnx_res, (out_info["w"]))
tabla_path = f"{OUTPATH}/{graph.name}{m}_tabla.json"
tabla_ir = pm.generate_tabla(graph,
shape_dict,
tabla_path)