def test_jax(self):
"""Test that a jax array is automatically converted into
a diagonal tensor"""
t = jnp.array([0.1, 0.2, 0.3])
res = fn.diag(t)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3]))
res = fn.diag(t, k=1)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3], k=1))
def test_array(self):
"""Test that a NumPy array is automatically converted into
a diagonal tensor"""
t = np.array([0.1, 0.2, 0.3])
res = fn.diag(t)
assert isinstance(res, np.ndarray)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3]))
res = fn.diag(t, k=1)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3], k=1))
def test_tensorflow(self):
"""Test that a tensorflow tensor is automatically converted into
a diagonal tensor"""
t = tf.Variable([0.1, 0.2, 0.3])
res = fn.diag(t)
assert isinstance(res, tf.Tensor)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3]))
res = fn.diag(t, k=1)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3], k=1))
def test_torch(self):
"""Test that a torch tensor is automatically converted into
a diagonal tensor"""
t = torch.tensor([0.1, 0.2, 0.3])
res = fn.diag(t)
assert isinstance(res, torch.Tensor)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3]))
res = fn.diag(t, k=1)
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.3], k=1))
def test_torch(self):
"""Test that a torch tensor is differentiable when using scatter addition"""
x = torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], requires_grad=True)
y = torch.tensor(0.56, requires_grad=True)
res = fn.scatter_element_add(x, [1, 2], y ** 2)
loss = res[1, 2]
assert isinstance(res, torch.Tensor)
assert fn.allclose(res.detach(), onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
loss.backward()
assert fn.allclose(x.grad, onp.array([[0, 0, 0], [0, 0, 1.0]]))
assert fn.allclose(y.grad, 2 * y)
def test_tensorflow(self):
"""Test that a TF tensor is differentiable when using scatter addition"""
x = tf.Variable([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
y = tf.Variable(0.56)
with tf.GradientTape() as tape:
res = fn.scatter_element_add(x, [1, 2], y ** 2)
loss = res[1, 2]
assert isinstance(res, tf.Tensor)
assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
grad = tape.gradient(loss, [x, y])
assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
assert fn.allclose(grad[1], 2 * y)
def test_array(self):
"""Test that a NumPy array is differentiable when using scatter addition"""
x = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], requires_grad=True)
y = np.array(0.56, requires_grad=True)
def cost(weights):
return fn.scatter_element_add(weights[0], [1, 2], weights[1] ** 2)
res = cost([x, y])
assert isinstance(res, np.ndarray)
assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
grad = qml.grad(lambda weights: cost(weights)[1, 2])([x, y])
assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
assert fn.allclose(grad[1], 2 * y)
def test_jax(self):
"""Test that a JAX array is differentiable when using scatter addition"""
x = jnp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
y = jnp.array(0.56)
def cost(weights):
return fn.scatter_element_add(weights[0], [1, 2], weights[1] ** 2)
res = cost([x, y])
assert isinstance(res, jax.interpreters.xla.DeviceArray)
assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
grad = jax.grad(lambda weights: cost(weights)[1, 2])([x, y])
assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
assert fn.allclose(grad[1], 2 * y)
def test_sequence(self, a, interface):
"""Test that a sequence is automatically converted into
a diagonal tensor"""
t = [0.1, 0.2, a]
res = fn.diag(t)
assert fn.get_interface(res) == interface
assert fn.allclose(res, onp.diag([0.1, 0.2, 0.5]))
def test_block_diag(tensors):
"""Tests for the block diagonal function"""
res = fn.block_diag(tensors)
expected = np.array(
[[1, 2, 0, 0, 0], [3, 4, 0, 0, 0], [0, 0, 1, 2, 0], [0, 0, -1, -6, 0], [0, 0, 0, 0, 5]]
)
assert fn.allclose(res, expected)
def test_multidimensional_indexing_along_axis(self, t):
"""Test that indexing with a sequence properly extracts
the elements from the specified tensor axis"""
indices = np.array([[0, 0], [1, 0]])
res = fn.take(t, indices, axis=1)
expected = np.array([[[[1, 2], [1, 2]], [[3, 4], [1, 2]]],
[[[5, 6], [5, 6]], [[0, -1], [5, 6]]]])
assert fn.allclose(res, expected)
def test_array_indexing_along_axis(self, t):
"""Test that indexing with a sequence properly extracts
the elements from the specified tensor axis"""
indices = [0, 1, -2]
res = fn.take(t, indices, axis=2)
expected = np.array([[[1, 2, 1], [3, 4, 3], [-1, 1, -1]],
[[5, 6, 5], [0, -1, 0], [2, 1, 2]]])
assert fn.allclose(res, expected)
def test_gather(tensor):
"""Tests for the gather function"""
indices = [1, 0]
res = fn.gather(tensor, indices)
expected = np.array([
[-1, -6, -3],
[ 1, 2, 3]
])
assert fn.allclose(res, expected)
def test_sum_axis_keepdims(self, t1):
"""Test that passing the axis argument allows for summing along
a specific axis, while keepdims avoids the summed dimensions from being removed"""
res = fn.sum_(t1, axis=(0, 2), keepdims=True)
# if tensorflow or pytorch, extract view of underlying data
if hasattr(res, "numpy"):
res = res.numpy()
assert fn.allclose(res, np.array([[[14], [6], [3]]]))
assert res.shape == (1, 3, 1)
def test_sum_axis(self, t1):
"""Test that passing the axis argument allows for summing along
a specific axis"""
res = fn.sum_(t1, axis=(0, 2))
# if tensorflow or pytorch, extract view of underlying data
if hasattr(res, "numpy"):
res = res.numpy()
assert fn.allclose(res, np.array([14, 6, 3]))
assert res.shape == (3, )
def test_concatenate_flattened_arrays(self, t1):
"""Concatenating arrays with axis=None will result in all arrays being pre-flattened"""
t2 = onp.array([5])
res = fn.concatenate([t1, t2], axis=None)
# if tensorflow or pytorch, extract view of underlying data
if hasattr(res, "numpy"):
res = res.numpy()
assert fn.allclose(res, np.array([1, 2, 5]))
assert list(res.shape) == [3]
def test_stack_axis(self, t1):
"""Test that passing the axis argument allows for stacking along
a different axis"""
t2 = onp.array([3, 4])
res = fn.stack([t1, t2], axis=1)
# if tensorflow or pytorch, extract view of underlying data
if hasattr(res, "numpy"):
res = res.numpy()
assert fn.allclose(res, np.array([[1, 3], [2, 4]]))
assert list(res.shape) == [2, 2]
def test_allclose(t1, t2):
"""Test that the allclose function works for a variety of inputs."""
res = fn.allclose(t1, t2)
if isinstance(t1, tf.Variable):
t1 = tf.convert_to_tensor(t1)
if isinstance(t2, tf.Variable):
t2 = tf.convert_to_tensor(t2)
expected = all(float(x) == float(y) for x, y in zip(t1, t2))
assert res == expected
def test_ones_like_explicit_dtype(self, t):
"""Test that the ones like function creates the correct
shape and type tensor."""
res = fn.ones_like(t, dtype=np.float16)
if isinstance(t, (list, tuple)):
t = onp.asarray(t)
assert res.shape == t.shape
assert fn.get_interface(res) == fn.get_interface(t)
assert fn.allclose(res, np.ones(t.shape))
# if tensorflow or pytorch, extract view of underlying data
if hasattr(res, "numpy"):
res = res.numpy()
t = t.numpy()
assert onp.asarray(res).dtype.type is np.float16
def test_sqrt(t):
"""Test that the square root function works for a variety
of input"""
res = fn.sqrt(t)
assert fn.allclose(res, [1, np.sqrt(2), np.sqrt(3)])
def test_where(t):
"""Test that the where function works as expected"""
res = fn.where(t < 0, 100 * fn.ones_like(t), t)
expected = np.array([[[1, 2], [3, 4], [100, 1]], [[5, 6], [0, 100], [2,
1]]])
assert fn.allclose(res, expected)
def test_torch(self):
"""Test that sum, called without the axis arguments, returns a scalar"""
t = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
res = fn.sum_(t)
assert isinstance(res, torch.Tensor)
assert fn.allclose(res, 2.1)
def test_multidimensional_indexing(self, t):
"""Test that indexing with a multi-dimensional sequence properly extracts
the elements from the flattened tensor"""
indices = [[0, 1], [3, 2]]
res = fn.take(t, indices)
assert fn.allclose(res, [[1, 2], [4, 3]])
def test_array_indexing(self, t):
"""Test that indexing with a sequence properly extracts
the elements from the flattened tensor"""
indices = [0, 2, 3, 6, -2]
res = fn.take(t, indices)
assert fn.allclose(res, [1, 3, 4, 5, 2])
def test_flattened_indexing(self, t):
"""Test that indexing without the axis argument
will flatten the tensor first"""
indices = 5
res = fn.take(t, indices)
assert fn.allclose(res, 1)
def test_jax(self):
"""Test that sum, called without the axis arguments, returns a scalar"""
t = jnp.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
res = fn.sum_(t)
assert fn.allclose(res, 2.1)
