def test_elementwise_trinary_out(self):
out = testing.shaped_random((30, 20, 40), cupy, self.dtype, seed=3)
desc_a = cutensor.create_tensor_descriptor(self.a)
desc_b = cutensor.create_tensor_descriptor(self.b)
desc_c = cutensor.create_tensor_descriptor(self.c)
d = cutensor.elementwise_trinary(self.alpha,
self.a,
desc_a,
self.mode_a,
self.beta,
self.b,
desc_b,
self.mode_b,
self.gamma,
self.c,
desc_c,
self.mode_c,
out=out)
assert d is out
testing.assert_allclose(self.alpha * self.a_transposed +
self.beta * self.b_transposed +
self.gamma * self.c,
d,
rtol=self.tol,
atol=self.tol)
def test_elementwise_trinary(self):
desc_a = cutensor.create_tensor_descriptor(self.a, ct.OP_SQRT)
desc_b = cutensor.create_tensor_descriptor(self.b, ct.OP_TANH)
desc_c = cutensor.create_tensor_descriptor(self.c, ct.OP_COS)
d = cutensor.elementwise_trinary(self.alpha,
self.a,
desc_a,
self.mode_a,
self.beta,
self.b,
desc_b,
self.mode_b,
self.gamma,
self.c,
desc_c,
self.mode_c,
op_AB=ct.OP_ADD,
op_ABC=ct.OP_MUL)
testing.assert_allclose((self.alpha * cupy.sqrt(self.a_transposed) +
self.beta * cupy.tanh(self.b_transposed)) *
self.gamma * cupy.cos(self.c),
d,
rtol=1e-6,
atol=1e-6)
def test_elementwise_trinary(self):
desc_a = cutensor.create_tensor_descriptor(self.a)
desc_b = cutensor.create_tensor_descriptor(self.b)
desc_c = cutensor.create_tensor_descriptor(self.c)
d = cutensor.elementwise_trinary(self.alpha, self.a, desc_a,
self.mode_a, self.beta, self.b,
desc_b, self.mode_b, self.gamma,
self.c, desc_c, self.mode_c)
assert d.dtype == self.dtype
testing.assert_allclose(self.alpha * self.a_transposed +
self.beta * self.b_transposed +
self.gamma * self.c_transposed,
d,
rtol=self.tol,
atol=self.tol)
def test_elementwise_trinary(self):
desc_a = cutensor.create_tensor_descriptor(self.a)
desc_b = cutensor.create_tensor_descriptor(self.b)
desc_c = cutensor.create_tensor_descriptor(self.c)
d = cutensor.elementwise_trinary(
self.alpha, self.a, desc_a, self.mode_a,
self.beta, self.b, desc_b, self.mode_b,
self.gamma, self.c, desc_c, self.mode_c
)
assert d.dtype == numpy.float32
testing.assert_allclose(
self.alpha.item() * self.a_transposed +
self.beta.item() * self.b_transposed +
self.gamma.item() * self.c_transposed,
d,
rtol=1e-6, atol=1e-6
)
c = cupy.random.random([extent[i] for i in mode_c])
a = a.astype(dtype)
b = b.astype(dtype)
c = c.astype(dtype)
desc_a = cutensor.create_tensor_descriptor(a)
desc_b = cutensor.create_tensor_descriptor(b)
desc_c = cutensor.create_tensor_descriptor(c)
alpha = 1.1
beta = 1.2
gamma = 1.3
# rehearsal
d = cutensor.elementwise_trinary(alpha, a, desc_a, mode_a,
beta, b, desc_b, mode_b,
gamma, c, desc_c, mode_c)
ev_start = stream.Event()
ev_end = stream.Event()
st = stream.Stream()
with st:
# measurement
ev_start.record()
d = cutensor.elementwise_trinary(alpha, a, desc_a, mode_a,
beta, b, desc_b, mode_b,
gamma, c, desc_c, mode_c)
ev_end.record()
st.synchronize()
elapsed_ms = stream.get_elapsed_time(ev_start, ev_end)
