def _verify(input_shape, diagonal_shape, dtype, k=0, align="RIGHT_LEFT"):
input = relay.var("input", relay.TensorType(input_shape, dtype))
diagonal = relay.var("diagonal", relay.TensorType(diagonal_shape, dtype))
out = relay.matrix_set_diag(input, diagonal, k, align)
in_type = run_infer_type(input)
out_type = run_infer_type(out)
assert in_type.checked_type == out_type.checked_type
func = relay.Function([input, diagonal], out)
input_np = np.random.randint(-100, 100, size=input_shape).astype(dtype)
diagonal_np = np.random.randint(-100, 100, size=diagonal_shape).astype(dtype)
out_np = tvm.topi.testing.matrix_set_diag(input_np, diagonal_np, k, align)
for target, dev in tvm.testing.enabled_targets():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, device=dev, target=target)
out_relay = intrp.evaluate(func)(input_np, diagonal_np)
tvm.testing.assert_allclose(out_relay.numpy(), out_np)
def _verify(input_shape, dtype):
diagonal_shape = list(input_shape[:-2])
diagonal_shape.append(min(input_shape[-2], input_shape[-1]))
input = relay.var("input", relay.TensorType(input_shape, dtype))
diagonal = relay.var("diagonal",
relay.TensorType(diagonal_shape, dtype))
out = relay.matrix_set_diag(input, diagonal)
in_type = run_infer_type(input)
out_type = run_infer_type(out)
assert in_type.checked_type == out_type.checked_type
func = relay.Function([input, diagonal], out)
input_np = np.random.randint(-100, 100, size=input_shape).astype(dtype)
diagonal_np = np.random.randint(-100, 100,
size=diagonal_shape).astype(dtype)
out_np = tvm.topi.testing.matrix_set_diag(input_np, diagonal_np)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
out_relay = intrp.evaluate(func)(input_np, diagonal_np)
tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
