Support for NumPy-style data types

.github/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### New features since last release
 
+* Python bindings now include a factory method which accepts a `dtype` parameter. [(#18)](https://github.com/XanaduAI/jet/pull/18)
+
 * Running `make build` from the `python` directory now creates a Python distribution package. [(#13)](https://github.com/XanaduAI/jet/pull/13)
 
 * A new intermediate representation (IR) is added, including a parser, IR representation program, and a Strawberry Fields interface. [(#11)](https://github.com/XanaduAI/jet/pull/11)

python/jet/__init__.py

@@ -1,12 +1,67 @@
+from typing import Union
+
+import numpy as np
+
 # The existence of a Python binding is proof of its intention to be exposed.
 from .bindings import *
 
-# By default, Python uses two 64-bit floating-point numbers to represent a
-# complex number. Altogether, this requires 128 bits of storage.
-Tensor = TensorC128
-TensorNetwork = TensorNetworkC128
-TensorNetworkFile = TensorNetworkFileC128
-TensorNetworkSerializer = TensorNetworkSerializerC128
+
+def Tensor(*args, **kwargs) -> Union[TensorC64, TensorC128]:
+    """Constructs a tensor with the specified data type. If a `dtype` keyword
+    argument is not provided, a TensorC128 instance will be returned.
+    """
+    dtype = kwargs.pop("dtype", "complex128")
+    if np.dtype(dtype) == np.complex64:
+        return TensorC64(*args, **kwargs)
+    elif np.dtype(dtype) == np.complex128:
+        return TensorC128(*args, **kwargs)
+    else:
+        raise TypeError(f"Data type '{dtype}' is not supported.")
+
+
+def TensorNetwork(*args, **kwargs) -> Union[TensorNetworkC64, TensorNetworkC128]:
+    """Constructs a tensor network with the specified data type. If a `dtype`
+    keyword argument is not provided, a TensorNetworkC128 instance will be
+    returned.
+    """
+    dtype = kwargs.pop("dtype", "complex128")
+    if np.dtype(dtype) == np.complex64:
+        return TensorNetworkC64(*args, **kwargs)
+    elif np.dtype(dtype) == np.complex128:
+        return TensorNetworkC128(*args, **kwargs)
+    else:
+        raise TypeError(f"Data type '{dtype}' is not supported.")
+
+
+def TensorNetworkFile(*args, **kwargs) -> Union[TensorNetworkFileC64, TensorNetworkFileC128]:
+    """Constructs a tensor network file with the specified data type. If a
+    `dtype` keyword argument is not provided, a TensorNetworkFileC128 instance
+    will be returned.
+    """
+    dtype = kwargs.pop("dtype", "complex128")
+    if np.dtype(dtype) == np.complex64:
+        return TensorNetworkFileC64(*args, **kwargs)
+    elif np.dtype(dtype) == np.complex128:
+        return TensorNetworkFileC128(*args, **kwargs)
+    else:
+        raise TypeError(f"Data type '{dtype}' is not supported.")
+
+
+def TensorNetworkSerializer(
+    *args, **kwargs
+) -> Union[TensorNetworkSerializerC64, TensorNetworkSerializerC128]:
+    """Constructs a tensor network serializer with the specified data type. If a
+    `dtype` keyword argument is not provided, a TensorNetworkSerializerC128
+    instance will be returned.
+    """
+    dtype = kwargs.pop("dtype", "complex128")
+    if np.dtype(dtype) == np.complex64:
+        return TensorNetworkSerializerC64(*args, **kwargs)
+    elif np.dtype(dtype) == np.complex128:
+        return TensorNetworkSerializerC128(*args, **kwargs)
+    else:
+        raise TypeError(f"Data type '{dtype}' is not supported.")
+
 
 # Grab the current Jet version from the C++ headers.
 __version__ = version()

python/setup.py

@@ -52,6 +52,7 @@ def build_extension(self, ext):
 
 requirements = [
     "lark-parser>=0.11.0",
+    "numpy",
     "StrawberryFields==0.18.0",
 ]
 

python/src/PathInfo.hpp

@@ -1,15 +1,17 @@
 #pragma once
 
+#include <string>
+
 #include <pybind11/pybind11.h>
 
 #include <Jet.hpp>
 
 namespace py = pybind11;
 
-template <class Tensor, class... Tensors>
+template <class T, class... Ts>
 void bind_constructors(py::class_<Jet::PathInfo> &c)
 {
-    c.def(py::init<const Jet::TensorNetwork<Tensor> &,
+    c.def(py::init<const Jet::TensorNetwork<Jet::Tensor<T>> &,
                    const Jet::PathInfo::path_t &>(),
           py::arg("tn"), py::arg("path") = Jet::PathInfo::path_t(), R"(
             Constructs a populated PathInfo for the given path
@@ -21,14 +23,19 @@ void bind_constructors(py::class_<Jet::PathInfo> &c)
                       the tensor network
             )");
 
-    if constexpr (sizeof...(Tensors) > 0) {
-        bind_constructors<Tensors...>(c);
+    if constexpr (sizeof...(Ts) > 0) {
+        bind_constructors<Ts...>(c);
     }
 }
 
-template <class... Tensors> void AddBindingsForPathInfo(py::module_ &m)
+/**
+ * @brief Adds Python bindings for the include/jet/PathInfo.hpp file.
+ *
+ * @tparam T Template parameter of the `Tensor` class.
+ * @param m Jet pybind11 module.
+ */
+template <class... Ts> void AddBindingsForPathInfo(py::module_ &m)
 {
-
     py::class_<Jet::PathStepInfo>(m, "PathStepInfo", R"(
         PathStepInfo represents the contraction metadata associated
         with a node in a TensorNetwork)")
@@ -86,6 +93,6 @@ template <class... Tensors> void AddBindingsForPathInfo(py::module_ &m)
             Computes the total memory required to contract the tensor
             network along this path.)");
 
-    // Add constructor bindings for each Tensor type
-    bind_constructors<Tensors...>(cls);
+    // Add constructor bindings for each Tensor data type.
+    bind_constructors<Ts...>(cls);
 }

python/src/Python.cpp

@@ -15,21 +15,16 @@ PYBIND11_MODULE(bindings, m)
     using c64_t = std::complex<float>;
     using c128_t = std::complex<double>;
 
-    using TensorC64 = Jet::Tensor<c64_t>;
-    using TensorC128 = Jet::Tensor<c128_t>;
+    AddBindingsForPathInfo<c64_t, c128_t>(m);
 
-    AddBindingsForTensor<c64_t>(m, "TensorC64");
-    AddBindingsForTensor<c128_t>(m, "TensorC128");
+    AddBindingsForTensor<c64_t>(m);
+    AddBindingsForTensor<c128_t>(m);
 
-    AddBindingsForTensorNetwork<TensorC64>(m, "TensorNetworkC64");
-    AddBindingsForTensorNetwork<TensorC128>(m, "TensorNetworkC128");
+    AddBindingsForTensorNetwork<c64_t>(m);
+    AddBindingsForTensorNetwork<c128_t>(m);
 
-    AddBindingsForTensorNetworkIO<TensorC64>(m, "TensorNetworkFileC64",
-                                             "TensorNetworkSerializerC64");
-    AddBindingsForTensorNetworkIO<TensorC128>(m, "TensorNetworkFileC128",
-                                              "TensorNetworkSerializerC128");
-
-    AddBindingsForPathInfo<TensorC64, TensorC128>(m);
+    AddBindingsForTensorNetworkIO<c64_t>(m);
+    AddBindingsForTensorNetworkIO<c128_t>(m);
 
     AddBindingsForVersion(m);
 }

python/src/Tensor.hpp

@@ -1,27 +1,30 @@
 #include <complex>
 #include <sstream>
+#include <string>
 
 #include <pybind11/complex.h>
 #include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include "Type.hpp"
 #include <Jet.hpp>
 
 namespace py = pybind11;
 
 /**
  * @brief Adds Python bindings for the include/jet/Tensor.hpp file.
  *
- * @tparam T Template parameter for Jet::Tensor.
+ * @tparam T Template parameter of the `Tensor` class.
  * @param m Jet pybind11 module.
- * @param name Name of the Tensor class binding.
  */
-template <class T> void AddBindingsForTensor(py::module_ &m, const char *name)
+template <class T> void AddBindingsForTensor(py::module_ &m)
 {
     using tensor_t = Jet::Tensor<T>;
 
-    py::class_<tensor_t>(m, name, R"(
+    const std::string class_name = "Tensor" + Type<T>::suffix;
+
+    py::class_<tensor_t>(m, class_name.c_str(), R"(
         This class represents an n-rank data structure of complex-valued data
         for tensor operations. We use the following conventions:
 
@@ -31,6 +34,13 @@ template <class T> void AddBindingsForTensor(py::module_ &m, const char *name)
             - "Shape" refers to the dimensions of a tensor; the number of
               dimensions is the rank of the tensor.
     )")
+        // Static properties
+        // ---------------------------------------------------------------------
+
+        .def_property_readonly_static(
+            "dtype", [](const py::object &) { return Type<T>::dtype; },
+            "Data type of this tensor.")
+
         // Static functions
         // ---------------------------------------------------------------------
 

python/src/TensorNetwork.hpp

@@ -1,38 +1,55 @@
 #pragma once
 
 #include <sstream>
+#include <string>
 #include <unordered_map>
 #include <vector>
 
 #include <pybind11/complex.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include "Type.hpp"
 #include <Jet.hpp>
 
 namespace py = pybind11;
 
-template <class Tensor>
-void AddBindingsForTensorNetwork(py::module_ &m, const char *name)
+/**
+ * @brief Adds Python bindings for the `TensorNetwork` class.
+ *
+ * @tparam T Template parameter of the `Tensor` class.
+ * @param m Jet pybind11 module.
+ */
+template <class T> void AddBindingsForTensorNetwork(py::module_ &m)
 {
+    using TensorNetwork = Jet::TensorNetwork<Jet::Tensor<T>>;
+    using node_id_t = typename Jet::TensorNetwork<Jet::Tensor<T>>::node_id_t;
+    using Node = typename Jet::TensorNetwork<Jet::Tensor<T>>::Node;
+    using Edge = typename Jet::TensorNetwork<Jet::Tensor<T>>::Edge;
 
-    using TensorNetwork = Jet::TensorNetwork<Tensor>;
-    using node_id_t = typename Jet::TensorNetwork<Tensor>::node_id_t;
-    using Node = typename Jet::TensorNetwork<Tensor>::Node;
-    using Edge = typename Jet::TensorNetwork<Tensor>::Edge;
+    const std::string class_name = "TensorNetwork" + Type<T>::suffix;
+    const std::string class_name_node = "TensorNetworkNode" + Type<T>::suffix;
+    const std::string class_name_edge = "TensorNetworkEdge" + Type<T>::suffix;
 
     auto cls =
-        py::class_<TensorNetwork>(m, name, R"(
+        py::class_<TensorNetwork>(m, class_name.c_str(), R"(
         TensorNetwork represents a tensor network)")
 
+            // Static properties
+            // -----------------------------------------------------------------
+
+            .def_property_readonly_static(
+                "dtype", [](const py::object &) { return Type<T>::dtype; },
+                "Data type of this tensor network.")
+
             // Constructors
-            // --------------------------------------------------------------------
+            // -----------------------------------------------------------------
 
             .def(py::init<>(), R"(
             Constructs an empty tensor network)")
 
             // Magic methods
-            // --------------------------------------------------------------------
+            // -----------------------------------------------------------------
 
             .def("__str__",
                  [](const TensorNetwork &tn) {
@@ -42,7 +59,7 @@ void AddBindingsForTensorNetwork(py::module_ &m, const char *name)
                  })
 
             // Properties
-            // --------------------------------------------------------------------
+            // -----------------------------------------------------------------
 
             .def_property_readonly("index_to_edge_map",
                                    &TensorNetwork::GetIndexToEdgeMap, R"(
@@ -77,7 +94,7 @@ void AddBindingsForTensorNetwork(py::module_ &m, const char *name)
                 The number of unique indices in this tensor network)")
 
             // Other
-            // --------------------------------------------------------------------
+            // -----------------------------------------------------------------
 
             .def("add_tensor", &TensorNetwork::AddTensor, py::arg("tensor"),
                  py::arg("tags") = std::vector<std::string>(), R"(
@@ -106,15 +123,15 @@ void AddBindingsForTensorNetwork(py::module_ &m, const char *name)
                  py::arg("path") = std::vector<std::pair<size_t, size_t>>(), R"(
                 Contracts a tensor network along an optionally-provided path)");
 
-    py::class_<Node>(cls, (std::string(name) + "Node").c_str())
+    py::class_<Node>(cls, class_name_node.c_str())
         .def_readonly("id", &Node::id)
         .def_readonly("name", &Node::name)
         .def_readonly("indices", &Node::indices)
         .def_readonly("tags", &Node::tags)
         .def_readonly("contracted", &Node::contracted)
         .def_readonly("tensor", &Node::tensor);
 
-    py::class_<Edge>(cls, (std::string(name) + "Edge").c_str())
+    py::class_<Edge>(cls, class_name_edge.c_str())
         .def_readonly("dim", &Edge::dim)
         .def_readonly("node_ids", &Edge::node_ids)
         .def("__eq__", &Edge::operator==);