Repository: google/TensorNetwork

Issue:
NumpyBackend.dtype returns wrong type
The return type of `NumpyBackend.dtype` is currently given by the type of the `dtype` argument at initialization, but it should always yield a `numpy.dtype` object:
```python 
be = NumpyBackend(dtype=np.float64)
be.dtype #this is a `type` object, i.e. a `np.float64`
```
The fix is 

```python 
class NumPyBackend(base_backend.BaseBackend):
  """See base_backend.BaseBackend for documentation."""

  def __init__(self, dtype: Optional[numpy.dtype] = None):
    super(NumPyBackend, self).__init__()
    self.np = numpy
    self.name = "numpy"
    self._dtype =dtype# -> self._dtype = numpy.dtype(dtype)

```


Discussion:
@mganahl , you're quite busy with other things, so I'd be happy to fix this and add tests if you're not on it yet.
That would be great, thanks!
btw, the jax backend needs also to be fixed
on it!

Resolving patch:
diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -14,7 +14,7 @@
 
 from typing import Any, Optional, Tuple
 from tensornetwork.backends.numpy import numpy_backend
-import numpy
+import numpy as np
 
 Tensor = Any
 
@@ -22,7 +22,7 @@
 class JaxBackend(numpy_backend.NumPyBackend):
   """See base_backend.BaseBackend for documentation."""
 
-  def __init__(self, dtype: Optional[numpy.dtype] = None):
+  def __init__(self, dtype: Optional[np.dtype] = None):
     super(JaxBackend, self).__init__()
     try:
       #pylint: disable=import-outside-toplevel
@@ -33,7 +33,7 @@ def __init__(self, dtype: Optional[numpy.dtype] = None):
     self.jax = jax
     self.np = self.jax.numpy
     self.name = "jax"
-    self._dtype = dtype
+    self._dtype = np.dtype(dtype) if dtype is not None else None
 
   def convert_to_tensor(self, tensor: Tensor) -> Tensor:
     result = self.jax.jit(lambda x: x)(tensor)
@@ -44,29 +44,34 @@ def convert_to_tensor(self, tensor: Tensor) -> Tensor:
     return result
 
   def concat(self, values: Tensor, axis: int) -> Tensor:
-    return numpy.concatenate(values, axis)
+    return np.concatenate(values, axis)
 
   def randn(self,
             shape: Tuple[int, ...],
-            dtype: Optional[numpy.dtype] = None,
+            dtype: Optional[np.dtype] = None,
             seed: Optional[int] = None) -> Tensor:
     if not seed:
-      seed = numpy.random.randint(0, 2**63)
+      seed = np.random.randint(0, 2**63)
     key = self.jax.random.PRNGKey(seed)
 
     if not dtype:
-      dtype = self.dtype if self.dtype is not None else numpy.float64
+      dtype = self.dtype if self.dtype is not None else np.dtype(np.float64)
 
     def cmplx_randn(complex_dtype, real_dtype):
+      real_dtype = np.dtype(real_dtype)
+      complex_dtype = np.dtype(complex_dtype)
+
       key_2 = self.jax.random.PRNGKey(seed + 1)
-      return self.jax.random.normal(
-          key, shape,
-          dtype=real_dtype) + complex_dtype(1j) * self.jax.random.normal(
-              key_2, shape, dtype=real_dtype)
 
-    if dtype is self.np.complex128:
+      real_part = self.jax.random.normal(key, shape, dtype=real_dtype)
+      complex_part = self.jax.random.normal(key_2, shape, dtype=real_dtype)
+      unit = (np.complex64(1j) if complex_dtype == np.dtype(np.complex64)
+              else np.complex128(1j))
+      return real_part + unit * complex_part
+
+    if dtype is np.dtype(self.np.complex128):
       return cmplx_randn(dtype, self.np.float64)
-    if dtype is self.np.complex64:
+    if dtype is np.dtype(self.np.complex64):
       return cmplx_randn(dtype, self.np.float32)
 
     return self.jax.random.normal(key, shape).astype(dtype)
diff --git a/tensornetwork/backends/numpy/numpy_backend.py b/tensornetwork/backends/numpy/numpy_backend.py
--- a/tensornetwork/backends/numpy/numpy_backend.py
+++ b/tensornetwork/backends/numpy/numpy_backend.py
@@ -26,7 +26,7 @@ def __init__(self, dtype: Optional[numpy.dtype] = None):
     super(NumPyBackend, self).__init__()
     self.np = numpy
     self.name = "numpy"
-    self._dtype = dtype
+    self._dtype = self.np.dtype(dtype) if dtype is not None else None
 
   def tensordot(self, a: Tensor, b: Tensor, axes: Sequence[Sequence[int]]):
     return self.np.tensordot(a, b, axes)
@@ -86,7 +86,7 @@ def convert_to_tensor(self, tensor: Tensor) -> Tensor:
       raise TypeError("Expected a `np.array` or scalar. Got {}".format(
           type(tensor)))
     result = self.np.asarray(tensor)
-    if self.dtype is not None and result.dtype != self.dtype:
+    if self.dtype is not None and self.np.dtype(result.dtype) != self.dtype:
       raise TypeError(
           "Backend '{}' cannot convert tensor of dtype {} to dtype {}".format(
               self.name, result.dtype, numpy.dtype(self.dtype)))
@@ -135,9 +135,11 @@ def randn(self,
       self.np.random.seed(seed)
 
     if not dtype:
-      dtype = self.dtype if self.dtype is not None else self.np.float64
+      dtype = (self.dtype if self.dtype is not None
+               else self.np.dtype(self.np.float64))
 
-    if (dtype is self.np.complex128) or (dtype is self.np.complex64):
+    if ((dtype is self.np.dtype(self.np.complex128))
+        or (dtype is self.np.dtype(self.np.complex64))):
       return self.np.random.randn(*shape).astype(
           dtype) + 1j * self.np.random.randn(*shape).astype(dtype)
     return self.np.random.randn(*shape).astype(dtype)


Test patch:
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -10,6 +10,17 @@
 np_dtypes = np_randn_dtypes + [np.complex64, np.complex128]
 
 
+@pytest.mark.parametrize("dtype", np_dtypes)
+def test_dtype(dtype):
+  backend = jax_backend.JaxBackend(dtype)
+  assert backend.dtype == np.dtype(dtype)
+  assert isinstance(backend.dtype, np.dtype)
+
+  backend = jax_backend.JaxBackend(np.dtype(dtype))
+  assert backend.dtype == np.dtype(dtype)
+  assert isinstance(backend.dtype, np.dtype)
+
+
 def test_tensordot():
   backend = jax_backend.JaxBackend()
   a = backend.convert_to_tensor(2 * np.ones((2, 3, 4)))
diff --git a/tensornetwork/backends/numpy/numpy_backend_test.py b/tensornetwork/backends/numpy/numpy_backend_test.py
--- a/tensornetwork/backends/numpy/numpy_backend_test.py
+++ b/tensornetwork/backends/numpy/numpy_backend_test.py
@@ -9,6 +9,16 @@
 np_dtypes = np_randn_dtypes + [np.complex64, np.complex128]
 
 
+@pytest.mark.parametrize("dtype", np_dtypes)
+def test_dtype(dtype):
+  backend = numpy_backend.NumPyBackend(dtype)
+  assert backend.dtype == np.dtype(dtype)
+  assert isinstance(backend.dtype, np.dtype)
+
+  backend = numpy_backend.NumPyBackend(np.dtype(dtype))
+  assert backend.dtype == np.dtype(dtype)
+  assert isinstance(backend.dtype, np.dtype)
+
 def test_tensordot():
   backend = numpy_backend.NumPyBackend()
   a = backend.convert_to_tensor(2 * np.ones((2, 3, 4)))