def verify_tensordot(sa, sb, axes):
a = np.random.uniform(low=-1.0, high=1.0, size=sa).astype(np.float32)
b = np.random.uniform(low=-1.0, high=1.0, size=sb).astype(np.float32)
c1 = np.tensordot(a, b, axes)
c2 = with_tvm(lambda A, B: topi.tensordot(A, B, axes), a, b)
tvm.testing.assert_allclose(c1, c2, rtol=1e-5, atol=1e-5)
def test_basic_operation():
np.random.seed(0)
shape = (10, 10)
x = te.var("x", dtype="float32")
k = te.reduce_axis((0, 10), name="k")
l = te.reduce_axis((0, 10), name="l")
A0 = te.placeholder(shape, name="A0")
A1 = te.placeholder(shape, name="A1")
zeros = np.zeros(shape)
B = te.compute(shape, lambda i, j: A0[i, j], name="B")
check_grad(B, [A0])
B = te.compute(shape, lambda i, j: A0[i, j] + A1[i, j], name="B")
check_grad(B, [A0, A1])
B = te.compute(shape, lambda i, j: A0[i, j] + A0[j, i], name="B")
check_grad(B, A0)
B = te.compute(shape, lambda i, j: te.floor(A0[i, j]), name="B")
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.ceil(A0[i, j]), name="B")
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.trunc(A0[i, j]), name="B")
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.round(A0[i, j]), name="B")
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: A0[i, j] + te.exp(A0[j, i]), name="B")
check_grad(B, A0)
B = te.compute(
shape,
lambda i, j: te.log(0.1 + te.abs(A0[i, j] + te.exp(A0[j, i]))),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sigmoid(A0[i, j] * A0[i, j] * A0[j, i]),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.tanh(A0[i, j] * A0[i, j] * A0[j, i]),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sqrt(A0[i, j] * A0[i, j] * A0[j, i]),
name="B")
check_grad(B, A0, data_range=(0.1, 10))
B = te.compute(shape,
lambda i, j: te.power(te.abs(A0[i, j]), A0[j, i]),
name="B")
check_grad(B, A0, data_range=(-4, 4))
B = te.compute(shape, lambda i, j: A0[i, j] * A0[j, i], name="B")
check_grad(B, A0)
B = te.compute((10, ),
lambda i: te.sum(A0[i, k] * A0[k, i], axis=k),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sum(A0[i, k] * A0[k, i] + 5, axis=k),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.max(A0[i, k] * A0[k, j] + 5, axis=k),
name="B")
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: A0[i, j] * (A1[j, i] + A0[j, i]),
name="B")
check_grad(B, [A0, A1])
B = te.compute(shape,
lambda i, j: te.sum(
A0[k, k] - A0[te.min(j + k, 9), j] * A0[i, k], axis=k),
name="B")
check_grad(B, A0)
def fcombine(x, y):
return x * y
def fidentity(t0):
return tvm.tir.const(1, t0)
prod = te.comm_reducer(fcombine, fidentity, name="prod")
B = te.compute((10, 10),
lambda i, j: prod(A0[i, k] + A0[k, i], axis=k),
name="B")
check_grad(B, A0)
X = te.placeholder((10, ), name="X")
A = te.compute((10, ), lambda i: X[i] + X[9 - i])
B = te.compute((10, ), lambda i: X[i] * X[9 - i])
Y = topi.tensordot(A, B, 1)
check_grad(Y, X)
X = te.placeholder((3, 3), name="X")
Y = topi.einsum("ii->i", (X))
check_grad(Y, X)