✨ Add DLPack tensor export functionality to AIGs

docs/machine_learning.md

@@ -51,7 +51,7 @@ dataset = [
 print(len(dataset), dataset[0].num_pis, dataset[0].num_gates)
 ```
 
-### NetworkX
+## NetworkX
 
 The [NetworkX](https://networkx.org/) adapter allows you to convert an AIG into a {py:class}`~networkx.DiGraph` object.
 This enables use of graph-analysis and graph-learning tooling that operates on NetworkX graphs. Once converted, you can
@@ -128,50 +128,136 @@ plt.margins(x=0.2)
 plt.show()
 ```
 
-## Truth Tables
+## DLPack Tensors
+
+For high-throughput ML pipelines, `aigverse` can export AIG objects directly as graph tensors (node
+attributes, edge indices, and edge attributes) utilizing the [DLPack](https://dmlc.github.io/dlpack/latest/) protocol via the {py:meth}`~aigverse.networks.Aig.to_graph_tensors` method. This allows
+zero-copy hand-off to modern tensor frameworks, such
+as [PyTorch](https://docs.pytorch.org/docs/stable/dlpack.html),
+[JAX](https://docs.jax.dev/en/latest/_autosummary/jax.dlpack.from_dlpack.html#jax.dlpack.from_dlpack),
+[TensorFlow](https://www.tensorflow.org/api_docs/python/tf/experimental/dlpack/from_dlpack), etc., through
+`from_dlpack`.
+
+Encoding and `dtype` mapping:
+
+- Edge encoding (`edge_attr`):
+  - `EdgeTensorEncoding.BINARY`: regular `0.0`, inverted `1.0`
+  - `EdgeTensorEncoding.SIGNED`: regular `+1.0`, inverted `-1.0`
+  - `EdgeTensorEncoding.ONE_HOT`: regular `[1.0, 0.0]`, inverted `[0.0, 1.0]`
+- Node encoding (`node_attr`):
+  - `NodeTensorEncoding.INTEGER`: `constant=0`, `pi=1`, `gate=2`, `po=3`
+  - `NodeTensorEncoding.ONE_HOT`: `[constant, pi, gate, po]`
+- Tensor `dtype`s:
+  - `edge_index`: `int64`
+  - `edge_attr`: `float32`
+  - `node_attr`: `float32`
+
+Tensor shapes follow the convention:
+
+$$
+\mathbf{edge\_index} \in \mathbb{Z}^{2 \times E}, \quad
+\mathbf{edge\_attr} \in \mathbb{R}^{E \times D_{\text{edge}}}, \quad
+\mathbf{node\_attr} \in \mathbb{R}^{N \times D_{\text{node}}}
+$$
+
+where $E$ is the number of edges and $N$ is the number of nodes. For edge features,
+$D_{\text{edge}} = 1$ for `BINARY` and `SIGNED`, and $D_{\text{edge}} = 2$ for `ONE_HOT`.
+The node feature width $D_{\text{node}}$ depends on the chosen node encoding and enabled optional features.
+
+:::{note}
+Current limitations of `to_graph_tensors`:
+
+- The export targets **combinational** networks. Sequential networks are not supported.
+- Exported tensors are backed by **CPU host memory** (NumPy-backed DLPack producer).
+- `torch.from_dlpack(...)` is zero-copy on CPU, but moving tensors to CUDA still allocates GPU memory and performs a host-to-device copy.
+- When `node_tts=True`, the export is restricted to at most 16 primary inputs due to the exponential growth of truth table size. This is a practical limit for ML applications, but it is not a fundamental limitation of the API.
+  :::
+
+```{code-cell} ipython3
+import torch
+
+from aigverse.networks import Aig, EdgeTensorEncoding, NodeTensorEncoding
+
+aig = Aig()
+a = aig.create_pi()
+b = aig.create_pi()
+g = aig.create_and(a, b)
+aig.create_po(g)
+
+dlpack_data = aig.to_graph_tensors(
+  node_encoding=NodeTensorEncoding.INTEGER,
+  edge_encoding=EdgeTensorEncoding.BINARY,
+    levels=True,
+    fanouts=True,
+    node_tts=False,
+)
 
-Truth tables can be easily converted to Python lists or [NumPy](https://numpy.org/) arrays, making them compatible with
-standard ML libraries such as [scikit-learn](https://scikit-learn.org/), [PyTorch](https://pytorch.org/), or
-[TensorFlow](https://www.tensorflow.org/). Since `TruthTable` objects are iterable, this conversion is direct and
-intuitive. You can use these arrays as labels or features in supervised learning tasks, or as part of a dataset for
-training and evaluating models.
+edge_index = torch.from_dlpack(dlpack_data["edge_index"])
+edge_attr = torch.from_dlpack(dlpack_data["edge_attr"])
+node_attr = torch.from_dlpack(dlpack_data["node_attr"])
+
+print(edge_index.shape, edge_attr.shape, node_attr.shape)
+```
+
+You can immediately construct sparse tensors in Python:
+
+```{code-cell} ipython3
+num_nodes = node_attr.shape[0]
+
+sparse_adj = torch.sparse_coo_tensor(
+    indices=edge_index,
+    values=edge_attr,
+    size=(num_nodes, num_nodes, edge_attr.shape[1]),
+)
+
+print(sparse_adj.shape)
+```
+
+The same export also works
+with [NumPy's DLPack consumer API](https://numpy.org/doc/stable/release/1.22.0-notes.html#add-nep-47-compatible-dlpack-support):
 
 ```{code-cell} ipython3
-from aigverse.utils import TruthTable
 import numpy as np
 
-# Create a simple truth table, e.g., a 3-input majority function
-tt = TruthTable(3)
-tt.create_from_hex_string("e8")
+edge_index_np = np.from_dlpack(aig.to_graph_tensors()["edge_index"])
+print(edge_index_np.shape)
+```
+
+## Truth Tables
 
-# Export to a list
-tt_list = list(tt)
-print(f"As list: {tt_list}")
+Truth tables are iterable, but for ML pipelines it is best to keep data in contiguous array/tensor form from the
+start. A practical pattern is:
 
-# Export to NumPy arrays of different types
-tt_np_bool = np.array(tt)
-print(f"As NumPy bool array:  {tt_np_bool}")
-tt_np_int = np.array(tt, dtype=np.int32)
-print(f"As NumPy int array:   {tt_np_int}")
-tt_np_float = np.array(tt, dtype=np.float64)
-print(f"As NumPy float array: {tt_np_float}")
+1. Materialize labels once as a NumPy array.
+2. Build the input matrix with vectorized NumPy operations.
+3. Convert to framework tensors (for example, [PyTorch](https://docs.pytorch.org/docs/stable/index.html)) without
+   copying.
 
+This keeps preprocessing fast and avoids Python-level loops in hot paths.
+
+```{code-cell} ipython3
+import numpy as np
+import torch
+
+from aigverse.utils import TruthTable
+
+# Create a simple truth table (3-input majority function)
+tt = TruthTable(3)
+tt.create_from_hex_string("e8")
 
-# These arrays can now be used as labels for an ML model.
-# For example, let's generate the corresponding feature matrix:
-def generate_inputs(num_vars):
-    inputs = []
-    for i in range(2**num_vars):
-        # Convert i to binary and pad with zeros
-        binary = bin(i)[2:].zfill(num_vars)
-        inputs.append([int(bit) for bit in binary])
-    return np.array(inputs)
+# Vectorized label extraction (shape: [2**num_vars])
+y_np = np.fromiter(tt, dtype=np.uint8)
 
+# Vectorized feature matrix generation (shape: [2**num_vars, num_vars])
+n = tt.num_vars()
+indices = np.array(range(2**n), dtype=np.uint32)[:, None]
+bit_positions = np.array(range(n - 1, -1, -1), dtype=np.uint32)
+X_np = ((indices >> bit_positions) & 1).astype(np.uint8)
 
-feature_matrix = generate_inputs(tt.num_vars())
-labels = tt_np_int  # Using the integer array as labels
+# Zero-copy bridge NumPy -> PyTorch
+X_torch = torch.from_numpy(X_np)
+y_torch = torch.from_numpy(y_np)
 
-print("\nFeature matrix (X) and labels (y) for ML:")
-print("X:\n", feature_matrix)
-print("y:\n", labels)
+print("NumPy shapes:", X_np.shape, y_np.shape)
+print("Torch shapes:", X_torch.shape, y_torch.shape)
 ```

pyproject.toml

@@ -42,7 +42,7 @@ requires-python = ">=3.10"
 dynamic = ["version"]
 
 [project.optional-dependencies]
-adapters = ["networkx>=3.0.0", "numpy>=1.23.0"]
+adapters = ["networkx>=3.0.0", "numpy>=1.23.0", "torch>=2.2.0"]
 
 [project.urls]
 Source = "https://github.com/marcelwa/aigverse"
@@ -224,6 +224,7 @@ build = [
 adapters = [
     "networkx>=3.0.0",
     "numpy>=1.23.0",
+    "torch>=2.2.0",
 ]
 docs = [
     { include-group = "adapters" },

python/aigverse/networks.pyi

@@ -4,13 +4,48 @@ This module includes network types, edge and index list utilities, and helper
 objects for structural manipulation.
 """
 
+import enum
 from collections.abc import Iterator, Sequence
 from typing import TYPE_CHECKING, NoReturn, overload
 
 if TYPE_CHECKING:
     import networkx as nx
     import numpy as np
 
+class NodeTensorEncoding(enum.Enum):
+    """Node encoding mode for exported graph tensors.
+
+    All node features use `float32`; only the categorical encoding scheme changes.
+    - `INTEGER`: Node classes are scalar labels in the first feature column.
+    - `ONE_HOT`: Node classes are one-hot vectors in `[constant, pi, gate, po]` order.
+    """
+
+    INTEGER = 0
+    """
+    Scalar node labels in the first feature column: 0=constant, 1=pi, 2=gate, 3=po.
+    """
+
+    ONE_HOT = 1
+    """One-hot node labels in [constant, pi, gate, po] order."""
+
+class EdgeTensorEncoding(enum.Enum):
+    """Edge encoding mode for exported graph tensors.
+
+    All edge features use `float32`; only the categorical encoding scheme changes.
+    - `BINARY`: Edge polarity is binary (regular=0.0, inverted=1.0).
+    - `SIGNED`: Edge polarity is signed (regular=+1.0, inverted=-1.0).
+    - `ONE_HOT`: Edge polarity is one-hot in `[regular, inverted]` order.
+    """
+
+    BINARY = 0
+    """Labels are encoded as 0.0 (regular) and 1.0 (inverted)."""
+
+    SIGNED = 1
+    """Labels are encoded as +1.0 (regular) and -1.0 (inverted)."""
+
+    ONE_HOT = 2
+    """One-hot edge labels in [regular, inverted] order."""
+
 class AigSignal:
     """Represents a signal in an AIG.
 
@@ -343,6 +378,41 @@ class Aig:
                     PO nodes (only for :class:`~aigverse.NamedAig`).
             """
 
+    def to_graph_tensors(
+        self,
+        node_encoding: NodeTensorEncoding = ...,
+        edge_encoding: EdgeTensorEncoding = ...,
+        *,
+        levels: bool = True,
+        fanouts: bool = False,
+        node_tts: bool = False,
+    ) -> dict:
+        """Exports graph tensors for machine-learning workflows.
+
+        Returns sparse graph topology and features as DLPack-compatible arrays.
+
+        Edge encoding mapping:
+            - ``BINARY``: regular=0.0, inverted=1.0
+            - ``SIGNED``: regular=+1.0, inverted=-1.0
+            - ``ONE_HOT``: regular=[1.0, 0.0], inverted=[0.0, 1.0]
+
+        Node encoding mapping:
+            - ``INTEGER``: constant=0, pi=1, gate=2, po=3
+            - ``ONE_HOT``: [constant, pi, gate, po]
+
+        Args:
+            node_encoding: Node encoding mode as :class:`~aigverse.networks.NodeTensorEncoding`.
+            edge_encoding: Edge encoding mode as :class:`~aigverse.networks.EdgeTensorEncoding`.
+            levels: Appends logic level as a node feature.
+            fanouts: Appends fanout size as a node feature.
+            node_tts: Appends simulated node/output truth-table bits.
+
+        Returns:
+            A dictionary with ``edge_index`` (shape ``(2, E)``, dtype ``int64``),
+            ``edge_attr`` (shape ``(E, D_edge)``, dtype ``float32``), and ``node_attr``
+            (shape ``(N, D_node)``, dtype ``float32``).
+        """
+
     def __len__(self) -> int:
         """Returns the number of nodes."""
 
@@ -636,6 +706,17 @@ class SequentialAig(Aig):
     def to_index_list(self) -> NoReturn:
         """Sequential networks cannot be encoded as combinational index lists."""
 
+    def to_graph_tensors(
+        self,
+        node_encoding: NodeTensorEncoding = ...,
+        edge_encoding: EdgeTensorEncoding = ...,
+        *,
+        levels: bool = True,
+        fanouts: bool = False,
+        node_tts: bool = False,
+    ) -> dict:
+        """Sequential networks cannot be exported as combinational graph tensors."""
+
     def __getstate__(self) -> NoReturn:
         """Sequential networks are not pickleable via combinational index-list state."""
 

src/aigverse/algorithms/refactoring.cpp

@@ -2,10 +2,9 @@
 // Created by marcel on 03.09.25.
 //
 
+#include "aigverse/algorithms/transform_helpers.hpp"
 #include "aigverse/types.hpp"
 
-#include "transform_helpers.hpp"
-
 #include <fmt/format.h>
 #include <mockturtle/algorithms/node_resynthesis/sop_factoring.hpp>
 #include <mockturtle/algorithms/refactoring.hpp>

src/aigverse/algorithms/resubstitution.cpp

@@ -2,10 +2,9 @@
 // Created by marcel on 03.09.25.
 //
 
+#include "aigverse/algorithms/transform_helpers.hpp"
 #include "aigverse/types.hpp"
 
-#include "transform_helpers.hpp"
-
 #include <mockturtle/algorithms/aig_resub.hpp>
 #include <mockturtle/algorithms/resubstitution.hpp>
 #include <nanobind/nanobind.h>

src/aigverse/networks/edge_list.hpp

@@ -8,6 +8,7 @@
 #include <fmt/ranges.h>  // NOLINT(misc-include-cleaner)
 #include <mockturtle/traits.hpp>
 
+#include <cstddef>
 #include <cstdint>
 #include <tuple>
 #include <vector>
@@ -129,12 +130,33 @@ struct edge_list
      */
     std::vector<edge<Ntk>> edges{};
 };
+/**
+ * @brief Converts a network to an edge list.
+ *
+ * @tparam Ntk Network type.
+ * @param ntk Network.
+ * @param regular_weight Weight assigned to non-inverted edges.
+ * @param inverted_weight Weight assigned to inverted edges.
+ * @return Edge list representation of the network.
+ */
 template <typename Ntk>
 [[nodiscard]] edge_list<typename Ntk::base_type> to_edge_list(const Ntk& ntk, const int64_t regular_weight = 0,
                                                               const int64_t inverted_weight = 1) noexcept
 {
     auto el = edge_list<typename Ntk::base_type>(ntk);
 
+    // estimate edge count for initial vector reservation
+    std::size_t edge_count =
+        (static_cast<std::size_t>(Ntk::max_fanin_size) * static_cast<std::size_t>(ntk.num_gates())) +
+        static_cast<std::size_t>(ntk.num_pos());
+
+    if constexpr (mockturtle::has_foreach_ri_v<Ntk> && mockturtle::has_ri_to_ro_v<Ntk>)
+    {
+        edge_count += static_cast<std::size_t>(ntk.num_registers());
+    }
+
+    el.edges.reserve(edge_count);
+
     // constants, primary inputs, and regular nodes
     ntk.foreach_node(
         [&ntk, regular_weight, inverted_weight, &el](const auto& n)