Refactor Python package layout and add example

Pysrc/tensorium/__init__.py

@@ -1 +1,26 @@
-from .tensorium import *
+"""Python interface to the Tensorium C++ bindings."""
+from __future__ import annotations
+
+from . import tensorium as _C
+
+from . import vectors
+from . import matrices
+from . import tensors
+
+# Re-export commonly used classes
+Vector = vectors.Vector
+Vectord = vectors.Vectord
+Matrix = matrices.Matrix
+Tensor2d = tensors.Tensor2d
+Tensor4d = tensors.Tensor4d
+
+__all__ = [
+    "Vector",
+    "Vectord",
+    "Matrix",
+    "Tensor2d",
+    "Tensor4d",
+    "vectors",
+    "matrices",
+    "tensors",
+]

Pysrc/tensorium/matrices.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from . import tensorium as _C
+
+Matrix = _C.Matrix
+
+
+def add(a: Matrix, b: Matrix) -> Matrix:
+    """Add two matrices."""
+    return _C.tns.add_mat(a, b)
+
+
+def sub(a: Matrix, b: Matrix) -> Matrix:
+    """Subtract matrix b from a."""
+    return _C.tns.sub_mat(a, b)
+
+
+def mul(a: Matrix, b: Matrix) -> Matrix:
+    """Matrix multiplication."""
+    return _C.tns.mul(a, b)
+
+__all__ = ["Matrix", "add", "sub", "mul"]

Pysrc/tensorium/matrices.pyi

@@ -0,0 +1,10 @@
+from .tensorium import Matrix
+
+
+def add(a: Matrix, b: Matrix) -> Matrix: ...
+
+def sub(a: Matrix, b: Matrix) -> Matrix: ...
+
+def mul(a: Matrix, b: Matrix) -> Matrix: ...
+
+__all__: list[str]

Pysrc/tensorium/tensorium.pyi

@@ -0,0 +1,38 @@
+from typing import List, Tuple
+
+class Vector:
+    def __init__(self, data: List[float]) -> None: ...
+    def __len__(self) -> int: ...
+    def __getitem__(self, idx: int) -> float: ...
+    def print(self) -> None: ...
+
+class Vectord(Vector):
+    ...
+
+class Matrix:
+    rows: int
+    cols: int
+    def __init__(self, rows: int, cols: int) -> None: ...
+    def fill(self, values: List[List[float]]) -> None: ...
+    def print(self) -> None: ...
+    def __getitem__(self, idx: Tuple[int, int]) -> float: ...
+    def __setitem__(self, idx: Tuple[int, int], val: float) -> None: ...
+
+class tns:
+    @staticmethod
+    def add_vec(a: Vector, b: Vector) -> Vector: ...
+    @staticmethod
+    def sub_vec(a: Vector, b: Vector) -> Vector: ...
+    @staticmethod
+    def scl_vec(a: Vector, f: float) -> Vector: ...
+    @staticmethod
+    def dot_vec(a: Vector, b: Vector) -> float: ...
+    @staticmethod
+    def add_mat(a: Matrix, b: Matrix) -> Matrix: ...
+    @staticmethod
+    def sub_mat(a: Matrix, b: Matrix) -> Matrix: ...
+    @staticmethod
+    def mul(a: Matrix, b: Matrix) -> Matrix: ...
+
+Tensor2d = None
+Tensor4d = None

Pysrc/tensorium/tensors.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from . import tensorium as _C
+
+Tensor2d = _C.tns.Tensor2d
+Tensor4d = _C.tns.Tensor4d
+
+
+def contract(t: Tensor4d, i: int, j: int):
+    """Contract a rank-4 tensor along axes (i, j)."""
+    return _C.tns.contract_tensor(t, i, j)
+
+__all__ = [
+    "Tensor2d",
+    "Tensor4d",
+    "contract",
+]

Pysrc/tensorium/tensors.pyi

@@ -0,0 +1,8 @@
+from .tensorium import tns
+
+Tensor2d = tns.Tensor2d
+Tensor4d = tns.Tensor4d
+
+def contract(t: Tensor4d, i: int, j: int): ...
+
+__all__: list[str]

Pysrc/tensorium/vectors.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+from typing import List
+
+from . import tensorium as _C
+
+Vector = _C.Vector
+Vectord = _C.Vectord
+
+
+def add(a: Vector, b: Vector) -> Vector:
+    """Add two vectors elementwise."""
+    return _C.tns.add_vec(a, b)
+
+
+def sub(a: Vector, b: Vector) -> Vector:
+    """Subtract vector b from a."""
+    return _C.tns.sub_vec(a, b)
+
+
+def scale(a: Vector, factor: float) -> Vector:
+    """Scale vector by a scalar."""
+    return _C.tns.scl_vec(a, factor)
+
+
+def dot(a: Vector, b: Vector) -> float:
+    """Dot product between two vectors."""
+    return _C.tns.dot_vec(a, b)
+
+
+__all__ = [
+    "Vector",
+    "Vectord",
+    "add",
+    "sub",
+    "scale",
+    "dot",
+]

Pysrc/tensorium/vectors.pyi

@@ -0,0 +1,13 @@
+from typing import List
+from .tensorium import Vector, Vectord
+
+
+def add(a: Vector, b: Vector) -> Vector: ...
+
+def sub(a: Vector, b: Vector) -> Vector: ...
+
+def scale(a: Vector, factor: float) -> Vector: ...
+
+def dot(a: Vector, b: Vector) -> float: ...
+
+__all__: list[str]

examples/basic_usage.py

@@ -0,0 +1,21 @@
+"""Basic usage example for tensorium Python bindings."""
+from tensorium.vectors import Vector, add as add_vec, dot
+from tensorium.matrices import Matrix, add as add_mat
+
+
+def main() -> None:
+    v1 = Vector([1.0, 2.0, 3.0])
+    v2 = Vector([4.0, 5.0, 6.0])
+    print("v1 + v2 =", add_vec(v1, v2))
+    print("dot(v1, v2) =", dot(v1, v2))
+
+    m1 = Matrix(2, 2)
+    m1.fill([[1.0, 2.0], [3.0, 4.0]])
+    m2 = Matrix(2, 2)
+    m2.fill([[5.0, 6.0], [7.0, 8.0]])
+    print("m1 + m2 =")
+    add_mat(m1, m2).print()
+
+
+if __name__ == "__main__":
+    main()