def test_multi_dim():
with pm.Node(name="elem4") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
x = pm.input("x", shape=(m, n))
w = pm.state("w", shape=(m, n))
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
w[i, j] = (w[i, j] * x[i, j])
m_ = 3
n_ = 4
x_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
w_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("w", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {}
for i in range(m_):
for j in range(n_):
input_info[f"w/w({i}, {j})"] = w_[i, j]
input_info[f"x/x({i}, {j})"] = x_[i, j]
fine_grained_eval = lowered_graph("w/w(2, 3)", input_info)
assert fine_grained_eval == np_result[2, 3]
def test_flatten_result_length():
with pm.Node(name="linear_reg") as graph:
m = pm.placeholder("m", type_modifier="param")
x = pm.placeholder("x", shape=(m), type_modifier="input")
y = pm.placeholder("y", type_modifier="input")
w = pm.placeholder("w", shape=(m), type_modifier="state")
mu = pm.placeholder("mu", default_val=1.0, type_modifier="param")
i = pm.index(0, (m - 1).set_name("m-1")).set_name("i")
h = pm.sum([i], (x[i] * w[i]).set_name("x*w"), name="h")
d = (h - y).set_name("h-y")
g = (d * x[i]).set_name("d*x")
w_ = (w[i] - (mu * g[i]).set_name("mu*g")).set_name(("w_out"))
shape_val_pass = NormalizeGraph({"m": 3})
count_pass = CountNodes()
flatten_pass = Lower({})
new_graph = shape_val_pass(graph)
flattened_g = flatten_pass(new_graph)
x = np.random.randint(0, 10, 10)
y = np.random.randint(0, 10, 1)[0]
w = np.random.randint(0, 10, 10)
orig_graph = count_pass(flattened_g)
def test_single_dim_op_slice():
with pm.Node(name="elem3") as graph:
m = pm.parameter(name="m")
x = pm.input("x", shape=m)
w = pm.state("w", shape=m)
i = pm.index(0, m - 1, name="i")
out = (w[i] * x[i])
w[i] = (out[i] - w[i])
m_ = 3
x_ = np.random.randint(0, 10, m_)
w_ = np.random.randint(0, 10, m_)
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_ - w_
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": 3})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("w", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
fine_grained_eval = lowered_graph("w/w(2,)", input_info)
assert fine_grained_eval == np_result[2]
def test_single_dim_norm():
with pm.Node(name="elem1") as graph:
m = pm.parameter("m")
x = pm.input("x", shape=m)
w = pm.state("w", shape=m)
i = pm.index(0, m - 1, name="i")
w[i] = (w[i] * x[i])
x_ = np.random.randint(0, 10, 3)
w_ = np.random.randint(0, 10, 3)
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": 3})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("w", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
fine_grained_eval = lowered_graph("w/w(1,)", input_info)
assert fine_grained_eval == np_result[1]
pb_path = f"{OUTPATH}/{graph.name}.srdfg"
pm.pb_store(lowered_graph, OUTPATH)
loaded_node = pm.pb_load(pb_path)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
fine_grained_eval = loaded_node("w/w(1,)", input_info)
assert fine_grained_eval == np_result[1]
def test_multi_dim_op_slice():
with pm.Node(name="elem2") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
mu = pm.parameter(name="mu", default=2.0)
x = pm.input(name="x", shape=(m, n))
w = pm.state(name="w", shape=(m, n))
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
out = (x[i, j] * w[i, j]).set_name("w_out")
w[i, j] = (mu * (out[i, j] - w[i, j]))
m_ = 3
n_ = 2
x_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
w_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
coarse_eval = graph("w", x=x_, w=w_)
np_result = (x_ * w_ - w_) * 2.0
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("w", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {}
for i in range(m_):
for j in range(n_):
input_info[f"w/w({i}, {j})"] = w_[i, j]
input_info[f"x/x({i}, {j})"] = x_[i, j]
fine_grained_eval = lowered_graph("w/w(2, 1)", input_info)
assert fine_grained_eval == np_result[2, 1]
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_shape_eval():
graph = linear_reg_graph_mg()
shape_val_pass = NormalizeGraph({"m": 3})
flatten_pass = Lower({})
new_graph = shape_val_pass(graph)
count_pass = CountNodes()
orig_graph = count_pass(new_graph)
assert new_graph["x"].shape == (3, )
def test_lower_group_op():
with pm.Node(name="linear_reg1") as graph:
m = pm.parameter(name="m")
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
i = pm.index(0, m - 1, name="i")
h = pm.sum([i], w[i] * x[i], name="h")
m_ = 3
n_ = 3
x_ = np.random.randint(0, 10, m_)
w_ = np.random.randint(0, 10, (m_))
np_result = np.sum(x_ * w_)
np.testing.assert_allclose(graph("h", {"w": w_, "x": x_}), np_result)
np.testing.assert_allclose(graph("h", w=w_, x=x_), np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("h", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
#
fine_grained_eval = lowered_graph("h/h(4,)", input_info)
assert fine_grained_eval == np_result
pb_path = f"{OUTPATH}/linear_reg1.srdfg"
pm.pb_store(lowered_graph, OUTPATH)
loaded_node = pm.pb_load(pb_path) #
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
loaded_res = loaded_node("h/h(4,)", input_info)
assert loaded_node.func_hash() == lowered_graph.func_hash()
assert loaded_res == np_result
def test_reco():
m_ = 3
n_ = 3
k_ = 2
shape_dict = {"m": n_, "k": k_, "n": n_}
graph, input_info, out_info, keys = reco(coarse=True, **shape_dict)
coarse_eval = graph(keys, input_info)
np.testing.assert_allclose(coarse_eval[0], out_info["w1"])
np.testing.assert_allclose(coarse_eval[1], out_info["w2"])
fgraph, input_info, out_info, keys = reco(coarse=False, **shape_dict)
lower_pass = Lower({})
lowered_graph = lower_pass(fgraph, {})
all_vals = lowered_graph(keys, input_info)
w1_elems = np.prod(out_info["w1"].shape)
w2_elems = np.prod(out_info["w2"].shape)
out1 = np.asarray(list(all_vals[0:w1_elems])).reshape(out_info["w1"].shape)
out2 = np.asarray(list(all_vals[w1_elems:])).reshape(out_info["w2"].shape)
np.testing.assert_allclose(out1, out_info["w1"])
np.testing.assert_allclose(out2, out_info["w2"])
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 = reco(coarse=False, **shape_dict)
loaded_res = loaded_node(keys, input_info)
lres1 = np.asarray(list(loaded_res[0:w1_elems])).reshape(out_info["w1"].shape)
lres2 = np.asarray(list(loaded_res[w1_elems:])).reshape(out_info["w2"].shape)
np.testing.assert_allclose(lres1, out_info["w1"])
np.testing.assert_allclose(lres2, out_info["w2"])
def test_flatten_reco():
with pm.Node(name="recommender") as graph:
m = pm.parameter("m")
n = pm.parameter("n")
k = pm.parameter("k")
x1 = pm.input("x1", shape=(k, ))
x2 = pm.input("x2", shape=(k, ))
r1 = pm.input("r1", shape=(m, ))
y1 = pm.input("y1", shape=(m, ))
r2 = pm.input("r2", shape=(n, ))
y2 = pm.input("y2", shape=(n, ))
w1 = pm.state("w1", shape=(m, k))
w2 = pm.state("w2", shape=(n, k))
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
l = pm.index(0, k - 1, name="l")
h1_sum = pm.sum([l], (w1[i, l] *
x2[l]).set_name("w1*x2")).set_name("h1_sum")
h1 = (h1_sum[i] * r1[i]).set_name("h1")
h2_sum = pm.sum([l], (w2[j, l] *
x1[l]).set_name("w2*x1")).set_name("h2_sum")
h2 = (h2_sum[j] * r2[j]).set_name("h2")
d1 = (h1[i] - y1[i]).set_name("d1")
d2 = (h2[j] - y2[j]).set_name("d2")
g1 = (d1[i] * x2[l]).set_name("g1")
g2 = (d2[j] * x1[l]).set_name("g2")
w1[i, l] = (w1[i, l] - g1[i, l])
w2[j, l] = (w2[j, l] - g2[j, l])
m_ = 3
n_ = 3
k_ = 2
input_info = {}
input_info["m"] = m_
input_info["n"] = n_
input_info["k"] = k_
input_info["w1"] = np.random.randint(1, 6, m_ * k_).reshape(m_, k_)
input_info["w2"] = np.random.randint(1, 6, n_ * k_).reshape(n_, k_)
input_info["x1"] = np.random.randint(1, 6, k_)
input_info["x2"] = np.random.randint(1, 6, k_)
input_info["r1"] = np.random.randint(0, 2, m_)
input_info["y1"] = np.random.randint(0, 6, m_)
input_info["r2"] = np.random.randint(0, 2, n_)
input_info["y2"] = np.random.randint(0, 6, n_)
out_info = numpy_reco(input_info)
shape_val_pass = NormalizeGraph({"m": m_, "n": n_, "k": k_})
flatten_pass = Lower({})
new_graph = shape_val_pass(graph)
test_res = new_graph(["w1", "w2"], input_info)
np.testing.assert_allclose(test_res[0], out_info["w1"])
np.testing.assert_allclose(test_res[1], out_info["w2"])
flattened_g = flatten_pass(new_graph)
input_info = {}
input_info["m"] = m_
input_info["n"] = n_
input_info["k"] = k_
input_info["w1"] = np.random.randint(1, 6, m_ * k_).reshape(m_, k_)
input_info["w2"] = np.random.randint(1, 6, n_ * k_).reshape(n_, k_)
input_info["x1"] = np.random.randint(1, 6, k_)
input_info["x2"] = np.random.randint(1, 6, k_)
input_info["r1"] = np.random.randint(0, 2, m_)
input_info["y1"] = np.random.randint(0, 6, m_)
input_info["r2"] = np.random.randint(0, 2, n_)
input_info["y2"] = np.random.randint(0, 6, n_)
new_out_info = numpy_reco(input_info)
pairs_w1 = list(
product(*tuple([np.arange(i) for i in input_info["w1"].shape])))
pairs_w2 = list(
product(*tuple([np.arange(i) for i in input_info["w2"].shape])))
w1_init = input_info["w1"]
for p in pairs_w1:
input_info[f"w1/w1({p[0]}, {p[1]})"] = input_info["w1"][p]
input_info.pop("w1")
w2_init = input_info["w2"]
for p in pairs_w2:
input_info[f"w2/w2({p[0]}, {p[1]})"] = input_info["w2"][p]
input_info.pop("w2")
for p in range(k_):
input_info[f"x1/x1({p},)"] = input_info["x1"][p]
input_info[f"x2/x2({p},)"] = input_info["x2"][p]
input_info.pop("x1")
input_info.pop("x2")
for p in range(m_):
input_info[f"r1/r1({p},)"] = input_info["r1"][p]
input_info[f"y1/y1({p},)"] = input_info["y1"][p]
input_info.pop("r1")
input_info.pop("y1")
for p in range(n_):
input_info[f"r2/r2({p},)"] = input_info["r2"][p]
input_info[f"y2/y2({p},)"] = input_info["y2"][p]
input_info.pop("r2")
input_info.pop("y2")
w1_keys = [f"w1/w1({p[0]}, {p[1]})" for p in pairs_w1]
w2_keys = [f"w2/w2({p[0]}, {p[1]})" for p in pairs_w2]
all_vals = flattened_g(w1_keys + w2_keys, input_info)
out1 = np.asarray(list(all_vals[0:6])).reshape(new_out_info["w2"].shape)
out2 = np.asarray(list(all_vals[6:])).reshape(new_out_info["w2"].shape)
np.testing.assert_allclose(
new_out_info["w1"],
np.asarray(list(all_vals[0:6])).reshape(new_out_info["w2"].shape))
np.testing.assert_allclose(
new_out_info["w2"],
np.asarray(list(all_vals[6:])).reshape(new_out_info["w2"].shape))