Repository: google/TensorNetwork

Issue:
Nuke Shell backend
We don't use it.


Discussion:

Resolving patch:
diff --git a/tensornetwork/backends/backend_factory.py b/tensornetwork/backends/backend_factory.py
--- a/tensornetwork/backends/backend_factory.py
+++ b/tensornetwork/backends/backend_factory.py
@@ -16,7 +16,6 @@
 from tensornetwork.backends.tensorflow import tensorflow_backend
 from tensornetwork.backends.numpy import numpy_backend
 from tensornetwork.backends.jax import jax_backend
-from tensornetwork.backends.shell import shell_backend
 from tensornetwork.backends.pytorch import pytorch_backend
 from tensornetwork.backends.symmetric import symmetric_backend
 from tensornetwork.backends import abstract_backend
@@ -24,7 +23,6 @@
     "tensorflow": tensorflow_backend.TensorFlowBackend,
     "numpy": numpy_backend.NumPyBackend,
     "jax": jax_backend.JaxBackend,
-    "shell": shell_backend.ShellBackend,
     "pytorch": pytorch_backend.PyTorchBackend,
     "symmetric": symmetric_backend.SymmetricBackend
 }
diff --git a/tensornetwork/backends/shell/__init__.py b/tensornetwork/backends/shell/__init__.py
deleted file mode 100644
--- a/tensornetwork/backends/shell/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2019 The TensorNetwork Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
diff --git a/tensornetwork/backends/shell/shell_backend.py b/tensornetwork/backends/shell/shell_backend.py
deleted file mode 100644
--- a/tensornetwork/backends/shell/shell_backend.py
+++ /dev/null
@@ -1,359 +0,0 @@
-# Copyright 2019 The TensorNetwork Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import functools
-import operator
-from tensornetwork.backends import abstract_backend
-#pylint: disable=line-too-long
-from typing import Optional, Sequence, Tuple, List, Any, Union, Type, Callable, Text
-import numpy as np
-
-
-class ShellTensor:
-
-  def __init__(self, shape: Tuple[int, ...], dtype=None):
-    self.shape = shape
-    self.dtype = dtype
-
-  def reshape(self, new_shape: Tuple[int, ...]):
-    self.shape = new_shape
-    return self
-
-
-Tensor = Any
-
-
-class ShellBackend(abstract_backend.AbstractBackend):
-  """See base_backend.BaseBackend for documentation."""
-
-  def __init__(self) -> None:
-    super(ShellBackend, self).__init__()
-    self.name = "shell"
-
-  def tensordot(self, a: Tensor, b: Tensor,
-                axes: Sequence[Sequence[int]]) -> Tensor:
-    # Does not work when axis < 0
-    gen_a = (x for i, x in enumerate(a.shape) if i not in axes[0])
-    gen_b = (x for i, x in enumerate(b.shape) if i not in axes[1])
-    return ShellTensor(tuple(self._concat_generators(gen_a, gen_b)))
-
-  def _concat_generators(self, *gen):
-    """Concatenates Python generators."""
-    for g in gen:
-      yield from g
-
-  def reshape(self, tensor: Tensor, shape: Sequence[int]) -> Tensor:
-    tensor = tensor.reshape(tuple(shape))
-    return tensor
-
-  def transpose(self, tensor: Tensor, perm: Sequence[int]) -> Tensor:
-    shape = tuple(tensor.shape[i] for i in perm)
-    tensor = tensor.reshape(tuple(shape))
-    return tensor
-
-  def svd(
-      self,
-      tensor: Tensor,
-      pivot_axis: int = -1,
-      max_singular_values: Optional[int] = None,
-      max_truncation_error: Optional[float] = None,
-      relative: Optional[bool] = False
-  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    if max_truncation_error is not None:
-      raise NotImplementedError("SVD with truncation shape cannot be "
-                                "calculated without explicit tensor values.")
-    left_dims = tensor.shape[:pivot_axis]
-    right_dims = tensor.shape[pivot_axis:]
-    dim_s0 = min(
-        functools.reduce(operator.mul, left_dims),
-        functools.reduce(operator.mul, right_dims))
-    if max_singular_values is not None:
-      dim_s = min(dim_s0, max_singular_values)
-    else:
-      dim_s = dim_s0
-
-    u = ShellTensor(left_dims + (dim_s,))
-    vh = ShellTensor((dim_s,) + right_dims)
-    s = ShellTensor((dim_s,))
-    s_rest = ShellTensor((dim_s0 - dim_s,))
-    return u, s, vh, s_rest
-
-  def qr(self, tensor: Tensor, pivot_axis: int = -1,
-         non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
-    left_dims = tensor.shape[:pivot_axis]
-    right_dims = tensor.shape[pivot_axis:]
-    center_dim = min(np.prod(left_dims), np.prod(right_dims))
-    q = ShellTensor(left_dims + (center_dim,))
-    r = ShellTensor((center_dim,) + right_dims)
-    return q, r
-
-  def rq(self, tensor: Tensor, pivot_axis: int = -1,
-         non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
-    left_dims = tensor.shape[:pivot_axis]
-    right_dims = tensor.shape[pivot_axis:]
-    center_dim = min(np.prod(left_dims), np.prod(right_dims))
-    q = ShellTensor(left_dims + (center_dim,))
-    r = ShellTensor((center_dim,) + right_dims)
-    return q, r
-
-  def shape_concat(self, values: Sequence[Tensor], axis: int) -> Tensor:
-    shape = values[0].shape
-    if axis < 0:
-      axis += len(shape)
-    concat_size = sum(v.shape[axis] for v in values)
-    new_shape = shape[:axis] + (concat_size,) + shape[axis + 1:]
-    return ShellTensor(new_shape)
-
-  def concat_shape(self, values) -> Sequence:
-    tuple_values = (tuple(v) for v in values)
-    return functools.reduce(operator.concat, tuple_values)
-
-  def shape_tensor(self, tensor: Tensor) -> Tuple:
-    return tensor.shape
-
-  def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]:
-    return tensor.shape
-
-  def sparse_shape(self, tensor):
-    return self.shape_tuple(tensor)
-
-  def shape_prod(self, values: Tensor) -> int:
-    # This is different from the BaseBackend prod!
-    # prod calculates the product of tensor elements and cannot implemented
-    # for shell tensors
-    # This returns the product of sizes instead
-    return self.shape_product(values.shape)
-
-  def shape_product(self, shape: Sequence[int]) -> int:
-    return functools.reduce(operator.mul, shape)
-
-  def sqrt(self, tensor: Tensor) -> Tensor:
-    return tensor
-
-  def convert_to_tensor(self, tensor: Any) -> Tensor:
-    shell_tensor = ShellTensor(tuple(tensor.shape))
-    return shell_tensor
-
-  def outer_product(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    return ShellTensor(tensor1.shape + tensor2.shape)
-
-  #pylint: disable=unused-argument
-  def einsum(self,
-             expression: str,
-             *tensors: Tensor,
-             optimize: bool = True) -> Tensor:
-    expr_list = expression.split(",")
-    expr_list[-1], res = expr_list[-1].split("->")
-    shape = tuple(self._find_char(expr_list, char, tensors) for char in res)
-    return ShellTensor(shape)
-
-  def _find_char(self, expr_list: List[str], char: str,
-                 tensors: Sequence[Tensor]) -> int:
-    """Finds character in einsum tensor expression.
-
-    Args:
-      expr_list: List with expression for input tensors in einsum.
-      char: One character string (letter) that corresponds to a specific
-        einsum component.
-
-    Returns:
-      size: Size of the axis that corresponds to this einsum expression
-        character.
-    """
-    for i, expr in enumerate(expr_list):
-      ind = expr.find(char)
-      if ind != -1:
-        return tensors[i].shape[ind]
-    raise ValueError("Einsum output expression contains letters not given"
-                     "in input.")
-
-  def norm(self, tensor: Tensor) -> Tensor:
-    return ShellTensor(())
-
-  def eye(self,
-          N: int,
-          dtype: Optional[Type[np.number]] = None,
-          M: Optional[int] = None) -> Tensor:
-    if not M:
-      M = N
-    return ShellTensor((N, M))
-
-  def ones(self,
-           shape: Tuple[int, ...],
-           dtype: Optional[Type[np.number]] = None) -> Tensor:
-    return ShellTensor(shape)
-
-  def zeros(self,
-            shape: Tuple[int, ...],
-            dtype: Optional[Type[np.number]] = None) -> Tensor:
-
-    return ShellTensor(shape)
-
-  def randn(self,
-            shape: Tuple[int, ...],
-            dtype: Optional[Type[np.number]] = None,
-            seed: Optional[int] = None) -> Tensor:
-    return ShellTensor(shape)
-
-  def random_uniform(self,
-                     shape: Tuple[int, ...],
-                     boundaries: Optional[Tuple[float, float]] = (0.0, 1.0),
-                     dtype: Optional[Type[np.number]] = None,
-                     seed: Optional[int] = None) -> Tensor:
-    return ShellTensor(shape)
-
-  def conj(self, tensor: Tensor) -> Tensor:
-    return tensor
-
-  def eigh(self, matrix: Tensor) -> Tuple[Tensor, Tensor]:
-    shape = matrix.shape
-    return ShellTensor((shape[0],)), ShellTensor(shape)
-
-  def eigs(self,
-           A: Callable,
-           args: Optional[List[Tensor]] = None,
-           initial_state: Optional[Tensor] = None,
-           shape: Optional[Tuple[int, ...]] = None,
-           dtype: Optional[Type[np.number]] = None,
-           num_krylov_vecs: Optional[int] = 200,
-           numeig: Optional[int] = 1,
-           tol: Optional[float] = 1E-8,
-           which: Optional[Text] = 'LR',
-           maxiter: Optional[int] = None) -> Tuple[Tensor, List]:
-    if args is None:
-      args = []
-
-    if num_krylov_vecs < numeig + 1:
-      raise ValueError('`num_krylov_vecs` >= `numeig+1` required!')
-
-    if initial_state is None:
-      if (shape is None) or (dtype is None):
-        raise ValueError("if no `initial_state` is passed, then `shape` and"
-                         "`dtype` have to be provided")
-      return [ShellTensor(tuple()) for _ in range(numeig)
-             ], [ShellTensor(shape) for _ in range(numeig)]
-    if not isinstance(initial_state, ShellTensor):
-      raise TypeError("Expected a `ShellTensor`. Got {}".format(
-          type(initial_state)))
-
-    return [ShellTensor(tuple()) for _ in range(numeig)
-           ], [ShellTensor(initial_state.shape) for _ in range(numeig)]
-
-  def eigsh_lanczos(self,
-                    A: Callable,
-                    args: Optional[List[Tensor]] = None,
-                    initial_state: Optional[Tensor] = None,
-                    shape: Optional[Tuple] = None,
-                    dtype: Optional[Type[np.number]] = None,
-                    num_krylov_vecs: int = 20,
-                    numeig: int = 1,
-                    tol: float = 1E-8,
-                    delta: float = 1E-8,
-                    ndiag: int = 20,
-                    reorthogonalize: bool = False) -> Tuple[Tensor, List]:
-    if args is None:
-      args = []
-    if num_krylov_vecs < numeig:
-      raise ValueError('`num_krylov_vecs` >= `numeig` required!')
-
-    if numeig > 1 and not reorthogonalize:
-      raise ValueError(
-          "Got numeig = {} > 1 and `reorthogonalize = False`. "
-          "Use `reorthogonalize=True` for `numeig > 1`".format(numeig))
-    if initial_state is None:
-      if (shape is None) or (dtype is None):
-        raise ValueError("if no `initial_state` is passed, then `shape` and"
-                         "`dtype` have to be provided")
-      return [ShellTensor(tuple()) for _ in range(numeig)
-             ], [ShellTensor(shape) for _ in range(numeig)]
-    if not isinstance(initial_state, ShellTensor):
-      raise TypeError("Expected a `ShellTensor`. Got {}".format(
-          type(initial_state)))
-
-    return [ShellTensor(tuple()) for _ in range(numeig)
-           ], [ShellTensor(initial_state.shape) for _ in range(numeig)]
-
-  def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    raise NotImplementedError("Shell tensor has not implemented addition( + )")
-
-  def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    raise NotImplementedError(
-        "Shell tensor has not implemented subtraction( - )")
-
-  def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    a = np.ones(tensor1.shape)
-    b = np.ones(tensor2.shape)
-    return ShellTensor((a * b).shape)
-
-  def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    raise NotImplementedError("Shell tensor has not implemented add( / )")
-
-  def index_update(self, tensor: Tensor, mask: Tensor,
-                   assignee: Tensor) -> Tensor:
-    return ShellTensor(tensor.shape)
-
-  def inv(self, matrix: Tensor) -> Tensor:
-    if len(matrix.shape) > 2:
-      raise ValueError(
-          "input to shell backend method `inv` has shape {}. Only matrices are supported."
-          .format(matrix.shape))
-    return ShellTensor(matrix.shape)
-
-  def broadcast_right_multiplication(self, tensor1: Tensor,
-                                     tensor2: Tensor) -> Tensor:
-    if len(tensor2.shape) != 1:
-      raise ValueError(
-          "only order-1 tensors are allowed for `tensor2`, found `tensor2.shape = {}`"
-          .format(tensor2.shape))
-
-    shape2 = tuple(tensor2.shape)
-    if len(shape2) < len(tensor1.shape):
-      shape2 = tuple([1] * (len(tensor1.shape) - len(shape2))) + shape2
-    shape = tuple([max([s1, s2]) for s1, s2 in zip(tensor1.shape, shape2)])
-    return ShellTensor(shape)
-
-  def broadcast_left_multiplication(self, tensor1: Tensor,
-                                    tensor2: Tensor) -> Tensor:
-    if len(tensor1.shape) != 1:
-      raise ValueError(
-          "only order-1 tensors are allowed for `tensor1`, found `tensor1.shape = {}`"
-          .format(tensor1.shape))
-
-    shape1 = tuple(tensor1.shape)
-    if len(shape1) < len(tensor2.shape):
-      shape1 = shape1 + tuple([1] * (len(tensor2.shape) - len(shape1)))
-    shape = tuple([max([s1, s2]) for s1, s2 in zip(tensor2.shape, shape1)])
-    return ShellTensor(shape)
-
-  def jit(self, fun: Callable, *args: List, **kwargs: dict) -> Callable:
-    return fun
-
-  def sum(self,
-          tensor: Tensor,
-          axis: Optional[Sequence[int]] = None,
-          keepdims: bool = False) -> Tensor:
-    if not keepdims:
-      newshape = np.delete(tensor.shape, axis)
-    else:
-      newshape = np.array(tensor.shape)
-      newshape[np.array(axis)] = 1
-    return ShellTensor(newshape)
-
-  def matmul(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    shape1 = np.array(tensor1.shape)[:-2]
-    shape2 = np.array(tensor2.shape)[:-2]
-    if not np.array_equal(shape1, shape2):
-      raise ValueError("shape mismatch for matmul")
-    new_shape = np.append(shape1, [tensor1.shape[-2], tensor2.shape[-1]])
-    return ShellTensor(new_shape)


Test patch:
diff --git a/tensornetwork/backends/shell/shell_backend_test.py b/tensornetwork/backends/shell/shell_backend_test.py
deleted file mode 100644
--- a/tensornetwork/backends/shell/shell_backend_test.py
+++ /dev/null
@@ -1,464 +0,0 @@
-# Copyright 2019 The TensorNetwork Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tests for tensornetwork.backends.shell.shell_backend."""
-
-import numpy as np
-import pytest
-from tensornetwork.backends.shell import shell_backend
-from tensornetwork.backends.numpy import numpy_backend
-
-
-def assertBackendsAgree(f, args):
-  np_result = getattr(numpy_backend.NumPyBackend(), f)(**args)
-  sh_result = getattr(shell_backend.ShellBackend(), f)(**args)
-  assert np_result.shape == sh_result.shape
-
-
-def test_shell_tensor_reshape():
-  shell_tensor = shell_backend.ShellTensor((2, 1), np.float64)
-  shell_tensor = shell_tensor.reshape((1, 2))
-  assert shell_tensor.shape == (1, 2)
-
-
-def test_tensordot():
-  args = {}
-  args["a"] = np.ones([3, 5, 2])
-  args["b"] = np.ones([2, 3])
-  args["axes"] = [[0, 2], [1, 0]]
-  assertBackendsAgree("tensordot", args)
-
-
-def test_reshape():
-  args = {"tensor": np.ones([3, 5, 2]), "shape": np.array([3, 10])}
-  assertBackendsAgree("reshape", args)
-
-
-def test_transpose():
-  args = {"tensor": np.ones([3, 5, 2]), "perm": [0, 2, 1]}
-  assertBackendsAgree("transpose", args)
-
-
-def test_svd():
-  tensor = np.ones([2, 3, 4, 5, 6])
-  np_res = numpy_backend.NumPyBackend().svd(tensor, 3)
-  sh_res = shell_backend.ShellBackend().svd(tensor, 3)
-  for x, y in zip(np_res, sh_res):
-    assert x.shape == y.shape
-
-
-def test_svd_raises_error():
-  tensor = np.ones([2, 3, 4, 5, 6])
-  with pytest.raises(NotImplementedError):
-    shell_backend.ShellBackend().svd(
-        tensor, 3, max_truncation_error=.1)
-
-
-def test_gmres_not_implemented():
-  backend = shell_backend.ShellBackend()
-  with pytest.raises(NotImplementedError):
-    backend.gmres(lambda x: x, np.ones((2)))
-
-
-def test_svd_with_max_values():
-  tensor = np.ones([2, 3, 4, 5, 6])
-  np_res = numpy_backend.NumPyBackend().svd(
-      tensor, 3, max_singular_values=5)
-  sh_res = shell_backend.ShellBackend().svd(
-      tensor, 3, max_singular_values=5)
-  for x, y in zip(np_res, sh_res):
-    assert x.shape == y.shape
-
-
-def test_qr():
-  tensor = np.ones([2, 3, 4, 5, 6])
-  np_res = numpy_backend.NumPyBackend().qr(tensor, 3)
-  sh_res = shell_backend.ShellBackend().qr(tensor, 3)
-  for x, y in zip(np_res, sh_res):
-    assert x.shape == y.shape
-
-
-def test_rq():
-  tensor = np.ones([2, 3, 4, 5, 6])
-  np_res = numpy_backend.NumPyBackend().rq(tensor, 3)
-  sh_res = shell_backend.ShellBackend().rq(tensor, 3)
-  for x, y in zip(np_res, sh_res):
-    assert x.shape == y.shape
-
-
-def test_shape_concat():
-  args = {
-      "values": [np.ones([3, 2, 5]),
-                 np.zeros([3, 2, 5]),
-                 np.ones([3, 3, 5])]
-  }
-  args["axis"] = 1
-  assertBackendsAgree("shape_concat", args)
-  args["axis"] = -2
-  assertBackendsAgree("shape_concat", args)
-
-
-def test_concat_shape():
-  shapes = [(5, 2), (3,), (4, 6)]
-  result = shell_backend.ShellBackend().concat_shape(shapes)
-  assert result == (5, 2, 3, 4, 6)
-
-
-def test_shape_tensor():
-  tensor = np.ones([3, 5, 2])
-  np_result = numpy_backend.NumPyBackend().shape_tensor(tensor)
-  sh_result = shell_backend.ShellBackend().shape_tensor(tensor)
-  assert np_result == sh_result
-
-
-def test_shape_tuple():
-  tensor = np.ones([3, 5, 2])
-  np_result = numpy_backend.NumPyBackend().shape_tuple(tensor)
-  sh_result = shell_backend.ShellBackend().shape_tuple(tensor)
-  assert np_result == sh_result
-
-
-def test_shape_prod():
-  result = shell_backend.ShellBackend().shape_prod(np.ones([3, 5, 2]))
-  assert result == 30
-
-
-def test_sqrt():
-  args = {"tensor": np.ones([3, 5, 2])}
-  assertBackendsAgree("sqrt", args)
-
-
-def test_convert_to_tensor():
-  args = {"tensor": np.ones([3, 5, 2])}
-  assertBackendsAgree("convert_to_tensor", args)
-
-
-def test_outer_product():
-  args = {"tensor1": np.ones([3, 5]), "tensor2": np.ones([4, 6])}
-  assertBackendsAgree("outer_product", args)
-
-
-def test_einsum():
-  expression = "ab,bc->ac"
-  tensor1, tensor2 = np.ones([5, 3]), np.ones([3, 6])
-  np_result = numpy_backend.NumPyBackend().einsum(expression, tensor1, tensor2)
-  sh_result = shell_backend.ShellBackend().einsum(expression, tensor1, tensor2)
-  assert np_result.shape == sh_result.shape
-
-
-def test_einsum_raises_error():
-  expression = "ab,bc->ad"
-  tensor1, tensor2 = np.ones([5, 3]), np.ones([3, 6])
-  with pytest.raises(ValueError):
-    shell_backend.ShellBackend().einsum(expression, tensor1, tensor2)
-
-
-def test_norm():
-  args = {"tensor": np.ones([3, 5])}
-  assertBackendsAgree("norm", args)
-
-
-def test_eye():
-  args = {"N": 10, "M": 8}
-  assertBackendsAgree("eye", args)
-
-
-def test_eye_without_M():
-  args = {"N": 10}
-  assertBackendsAgree("eye", args)
-
-
-def test_zeros():
-  args = {"shape": (10, 4)}
-  assertBackendsAgree("zeros", args)
-
-
-def test_conj():
-  args = {"tensor": np.ones([3, 5])}
-  assertBackendsAgree("conj", args)
-
-
-def test_ones():
-  args = {"shape": (10, 4)}
-  assertBackendsAgree("ones", args)
-
-
-def test_randn():
-  args = {"shape": (10, 4)}
-  assertBackendsAgree("randn", args)
-
-
-def test_random_uniform():
-  args = {"shape": (10, 4)}
-  assertBackendsAgree("random_uniform", args)
-
-
-def test_eigsh_lanczos_1():
-  backend = shell_backend.ShellBackend()
-  D = 16
-  init = backend.randn((D,))
-  eigvals, eigvecs = backend.eigsh_lanczos(
-      lambda x: x, initial_state=init, numeig=3, reorthogonalize=True)
-  for n, ev in enumerate(eigvals):
-    assert eigvecs[n].shape == (D,)
-    assert ev.shape == tuple()
-
-
-def test_eigsh_lanczos_shape():
-  backend = shell_backend.ShellBackend()
-  D = 16
-
-  def mv(x):
-    return x
-
-  eigvals, eigvecs = backend.eigsh_lanczos(
-      mv, shape=(D,), dtype=np.float64, numeig=3, reorthogonalize=True)
-
-  for n, ev in enumerate(eigvals):
-    assert eigvecs[n].shape == (D,)
-    assert ev.shape == tuple()
-
-
-def test_eigsh_lanczos_init_shape():
-  backend = shell_backend.ShellBackend()
-  D = 16
-  init = backend.randn((D,))
-
-  def mv(x):
-    return x
-
-  eigvals, eigvecs = backend.eigsh_lanczos(
-      mv, numeig=3, initial_state=init, reorthogonalize=True)
-  for n, ev in enumerate(eigvals):
-    assert eigvecs[n].shape == (D,)
-    assert ev.shape == tuple()
-
-
-def test_eigsh_lanczos_raises():
-  backend = shell_backend.ShellBackend()
-  with pytest.raises(ValueError):
-    backend.eigsh_lanczos(lambda x: x, numeig=10, num_krylov_vecs=9)
-  with pytest.raises(ValueError):
-    backend.eigsh_lanczos(lambda x: x, numeig=2, reorthogonalize=False)
-  with pytest.raises(ValueError):
-    backend.eigsh_lanczos(lambda x: x, shape=(10,), dtype=None)
-  with pytest.raises(ValueError):
-    backend.eigsh_lanczos(lambda x: x, shape=None, dtype=np.float64)
-  with pytest.raises(ValueError):
-    backend.eigsh_lanczos(lambda x: x)
-  with pytest.raises(TypeError):
-    backend.eigsh_lanczos(lambda x: x, initial_state=[1, 2, 3])
-
-
-@pytest.mark.parametrize("a, b", [
-    pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3))),
-    pytest.param(2. * np.ones(()), np.ones((1, 2, 3))),
-])
-def test_multiply(a, b):
-  args = {"tensor1": a, "tensor2": b}
-  assertBackendsAgree("multiply", args)
-
-
-def test_eigh():
-  matrix = np.ones([3, 3])
-  vals, vecs = shell_backend.ShellBackend().eigh(matrix)
-  assert vals.shape == (3,)
-  assert vecs.shape == (3, 3)
-
-
-def test_eigs():
-  backend = shell_backend.ShellBackend()
-  init = shell_backend.ShellTensor((2,), np.float64)
-  eta, v = backend.eigs(lambda x: x, initial_state=init, numeig=2)
-  assert len(eta) == 2
-  for n in range(len(eta)):
-    assert v[n].shape == (2,)
-
-  def mv(x):
-    return x
-
-  eta, v = backend.eigs(mv, shape=(2,), dtype=np.float64, numeig=2)
-  assert len(eta) == 2
-  for n in range(len(eta)):
-    assert v[n].shape == (2,)
-
-
-def test_eigs_initial_state_shape():
-  backend = shell_backend.ShellBackend()
-
-  def mv(x):
-    return x
-
-  eta, v = backend.eigs(mv, initial_state=backend.randn((2,)))
-  assert len(eta) == 1
-  for n in range(len(eta)):
-    assert v[n].shape == (2,)
-
-
-def test_eigs_raises():
-  backend = shell_backend.ShellBackend()
-  with pytest.raises(ValueError):
-    backend.eigs(lambda x: x, numeig=10, num_krylov_vecs=10)
-  with pytest.raises(
-      ValueError,
-      match="if no `initial_state` is passed, then `shape` and"
-      "`dtype` have to be provided"):
-    backend.eigs(lambda x: x, shape=(10,), dtype=None)
-  with pytest.raises(
-      ValueError,
-      match="if no `initial_state` is passed, then `shape` and"
-      "`dtype` have to be provided"):
-    backend.eigs(lambda x: x, shape=None, dtype=np.float64)
-  with pytest.raises(
-      ValueError,
-      match="if no `initial_state` is passed, then `shape` and"
-      "`dtype` have to be provided"):
-    backend.eigs(lambda x: x)
-  with pytest.raises(TypeError):
-    backend.eigs(lambda x: x, initial_state=[1, 2, 3])
-
-
-def index_update():
-  backend = shell_backend.ShellBackend()
-  tensor_1 = np.ones([2, 3, 4])
-  tensor_2 = backend.index_update(tensor_1, tensor_1 > 0.1, 0)
-  assert tensor_1.shape == tensor_2.shape
-
-
-def test_matrix_inv():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4), seed=10)
-  inverse = backend.inv(matrix)
-  assert inverse.shape == matrix.shape
-
-
-def test_matrix_inv_raises():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4, 4), seed=10)
-  with pytest.raises(ValueError):
-    backend.inv(matrix)
-
-
-def test_broadcast_right_multiplication():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((2, 4, 3))
-  tensor2 = backend.randn((3,))
-  out = backend.broadcast_right_multiplication(tensor1, tensor2)
-  np.testing.assert_allclose(out.shape, [2, 4, 3])
-
-
-def test_broadcast_right_multiplication_reverse_order():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((3,))
-  tensor2 = backend.randn((3,))
-  out = backend.broadcast_right_multiplication(tensor1, tensor2)
-  np.testing.assert_allclose(out.shape, [3])
-
-
-def test_broadcast_right_multiplication_raises():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((2, 4, 3))
-  tensor2 = backend.randn((3, 3))
-  with pytest.raises(ValueError):
-    backend.broadcast_right_multiplication(tensor1, tensor2)
-
-
-def test_broadcast_left_multiplication():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((3,))
-  tensor2 = backend.randn((3, 4, 2))
-  out = backend.broadcast_left_multiplication(tensor1, tensor2)
-  np.testing.assert_allclose(out.shape, [3, 4, 2])
-
-
-def test_broadcast_left_multiplication_reverse_order():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((3,))
-  tensor2 = backend.randn((3,))
-  out = backend.broadcast_left_multiplication(tensor1, tensor2)
-  np.testing.assert_allclose(out.shape, [3])
-
-
-def test_broadcast_left_multiplication_raises():
-  backend = shell_backend.ShellBackend()
-  tensor1 = backend.randn((3, 3))
-  tensor2 = backend.randn((3, 4, 2))
-  with pytest.raises(ValueError):
-    backend.broadcast_left_multiplication(tensor1, tensor2)
-
-
-def test_sparse_shape():
-  backend = shell_backend.ShellBackend()
-  tensor = backend.randn((2, 3, 4), seed=10)
-  np.testing.assert_allclose(backend.sparse_shape(tensor), tensor.shape)
-
-
-def test_addition():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4, 4), seed=10)
-  with pytest.raises(NotImplementedError):
-    backend.addition(matrix, matrix)
-
-
-def test_subtraction():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4, 4), seed=10)
-  with pytest.raises(NotImplementedError):
-    backend.subtraction(matrix, matrix)
-
-
-def test_divide():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4, 4), seed=10)
-  with pytest.raises(NotImplementedError):
-    backend.divide(matrix, matrix)
-
-
-def test_index_update():
-  backend = shell_backend.ShellBackend()
-  matrix = backend.randn((4, 4, 4), seed=10)
-  actual = backend.index_update(matrix, matrix, matrix)
-  assert isinstance(actual, shell_backend.ShellTensor)
-  assert actual.shape == (4, 4, 4)
-
-
-def test_sum():
-  np.random.seed(10)
-  backend = shell_backend.ShellBackend()
-  a = backend.randn((2, 3, 4), seed=10)
-  actual = backend.sum(a, axis=(1, 2))
-  np.testing.assert_allclose(actual.shape, [
-      2,
-  ])
-
-  actual = backend.sum(a, axis=(1, 2), keepdims=True)
-  np.testing.assert_allclose(actual.shape, [2, 1, 1])
-
-
-def test_matmul():
-  np.random.seed(10)
-  backend = shell_backend.ShellBackend()
-  a = backend.randn((10, 2, 3), seed=10)
-  b = backend.randn((10, 3, 4), seed=10)
-  actual = backend.matmul(a, b)
-  np.testing.assert_allclose(actual.shape, [10, 2, 4])
-
-
-@pytest.mark.parametrize("pivot_axis", [-1, 1, 2])
-def test_pivot(pivot_axis):
-  shape = (4, 3, 2, 8)
-  backend = numpy_backend.NumPyBackend()
-  pivot_shape = (np.prod(shape[:pivot_axis]), np.prod(shape[pivot_axis:]))
-  tensor = backend.randn(shape, dtype=np.float64)
-  actual = backend.pivot(tensor, pivot_axis=pivot_axis)
-  np.testing.assert_allclose(actual.shape, pivot_shape)