def get_elem(a, b, **kwargs):
if len(a.shape) == len(b.shape):
return pm.matmul(a, b, **kwargs)
elif len(a.shape) > len(b.shape):
return pm.rvmatmul(a, b, **kwargs)
else:
return pm.lvmatmul(a, b, **kwargs)
def test_matmul(in_shape, w_shape):
x = np.random.randint(0, 30, np.prod(in_shape)).reshape(in_shape)
w = np.random.randint(0, 30, np.prod(w_shape)).reshape(w_shape)
if in_shape[-1] == w_shape[-1]:
o_np = [email protected]
else:
assert in_shape[-1] == w_shape[0]
o_np = x @ w
with pm.Node(name="mmul") as graph:
x_pm = pm.input(name="x", shape=in_shape)
w_pm = pm.state(name="w", shape=w_shape)
o_pm = pm.output(name="o", shape=o_np.shape)
pm.matmul(x_pm, w_pm, o_pm)
in_dict = {"x": x, "w": w}
res = graph("o", in_dict)
np.testing.assert_allclose(o_np, res)
def get_matmul(a, b, out=None, **kwargs):
if not out:
out = pm.output(shape=kwargs['shape'], name=kwargs['name'])
pm.matmul(a, b, out)
return out