add affine transform util (#322)

* add affine transform util

* add test

* fix pyright

* fix faulty test

* fix dim and num results

* formatting

compiler/ir/autoflow/__init__.py

@@ -0,0 +1 @@
+from .affine_transform import *

compiler/ir/autoflow/affine_transform.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import Self
+
+
+@dataclass(frozen=True)
+class AffineTransform:
+    """
+    An affine transform mirroring the functionality of xDSLs and MLIRs
+    AffineMap, but represented in matrix form to make life much easier.
+    This is possible if you don't have to support floordiv/ceildiv operations.
+    """
+
+    A: npt.NDArray[np.int_]  # Transformation matrix
+    b: npt.NDArray[np.int_]  # Translation vector
+
+    def __post_init__(self):
+        # Validate dimensions
+        if self.A.ndim != 2:
+            raise ValueError("Matrix A must be 2-dimensional.")
+        if self.b.ndim != 1:
+            raise ValueError("Vector b must be 1-dimensional.")
+        if self.A.shape[0] != self.b.shape[0]:
+            raise ValueError("Matrix A and vector b must have compatible dimensions.")
+
+    @property
+    def num_dims(self) -> int:
+        return self.A.shape[1]
+
+    @property
+    def num_results(self) -> int:
+        return self.A.shape[0]
+
+    def eval(self, x: npt.NDArray[np.int_]) -> npt.NDArray[np.int_]:
+        """
+        Apply the affine transformation to a vector or a set of vectors.
+        """
+        if x.ndim == 1:  # Single vector
+            if x.shape[0] != self.A.shape[1]:
+                raise ValueError(
+                    "Input vector x must have a dimension matching the number of columns in A."
+                )
+            return self.A @ x + self.b
+        elif x.ndim == 2:  # Batch of vectors
+            if x.shape[1] != self.A.shape[1]:
+                raise ValueError(
+                    "Input vectors in batch must have a dimension matching the number of columns in A."
+                )
+            return (self.A @ x.T).T + self.b
+        else:
+            raise ValueError("Input x must be 1D (vector) or 2D (batch of vectors).")
+
+    def compose(self, other: Self) -> Self:
+        """
+        Combine this affine transformation with another.
+        The result represents the application of `other` followed by `self`.
+        """
+        if self.A.shape[1] != other.A.shape[0]:
+            raise ValueError(
+                "Matrix dimensions of the transformations do not align for composition."
+            )
+        new_A = self.A @ other.A
+        new_b = self.A @ other.b + self.b
+        return type(self)(new_A, new_b)
+
+    def __str__(self):
+        return f"AffineTransform(A=\n{self.A},\nb={self.b})"

tests/ir/autoflow/test_affine_transform.py

@@ -0,0 +1,97 @@
+import numpy as np
+import pytest
+
+from compiler.ir.autoflow import AffineTransform
+
+
+def test_affine_transform_initialization_valid():
+    A = np.array([[1, 2, 3], [4, 5, 6]])
+    b = np.array([7, 8])
+    transform = AffineTransform(A, b)
+    assert np.array_equal(transform.A, A)
+    assert np.array_equal(transform.b, b)
+    assert transform.num_dims == 3
+    assert transform.num_results == 2
+
+
+def test_affine_transform_initialization_invalid_dimensions():
+    A = np.array([[1, 2], [3, 4]])
+    b = np.array([5])  # Incompatible size
+    with pytest.raises(
+        ValueError, match="Matrix A and vector b must have compatible dimensions."
+    ):
+        AffineTransform(A, b)
+
+
+def test_affine_transform_eval_single_vector():
+    A = np.array([[1, 0], [0, 1]])
+    b = np.array([1, 2])
+    transform = AffineTransform(A, b)
+    x = np.array([3, 4])
+    result = transform.eval(x)
+    expected = np.array([4, 6])
+    assert np.array_equal(result, expected)
+
+
+def test_affine_transform_eval_batch_of_vectors():
+    A = np.array([[1, 0], [0, 1]])
+    b = np.array([1, 2])
+    transform = AffineTransform(A, b)
+    x_batch = np.array([[3, 4], [5, 6]])
+    result = transform.eval(x_batch)
+    expected = np.array([[4, 6], [6, 8]])
+    assert np.array_equal(result, expected)
+
+
+def test_affine_transform_eval_invalid_vector_dimension():
+    A = np.array([[1, 0], [0, 1]])
+    b = np.array([1, 2])
+    transform = AffineTransform(A, b)
+    x = np.array([3])  # Incompatible dimension
+    with pytest.raises(
+        ValueError,
+        match="Input vector x must have a dimension matching the number of columns in A.",
+    ):
+        transform.eval(x)
+
+
+def test_affine_transform_compose():
+    A1 = np.array([[1, 2], [3, 4]])
+    b1 = np.array([5, 6])
+    transform1 = AffineTransform(A1, b1)
+
+    A2 = np.array([[0, 1], [1, 0]])
+    b2 = np.array([7, 8])
+    transform2 = AffineTransform(A2, b2)
+
+    composed = transform1.compose(transform2)
+
+    expected_A = A1 @ A2
+    expected_b = A1 @ b2 + b1
+
+    assert np.array_equal(composed.A, expected_A)
+    assert np.array_equal(composed.b, expected_b)
+
+
+def test_affine_transform_compose_invalid_dimensions():
+    A1 = np.array([[1, 2], [3, 4]])
+    b1 = np.array([5, 6])
+    transform1 = AffineTransform(A1, b1)
+
+    A2 = np.array([[1, 0, 0], [0, 1, 0]])
+    b2 = np.array([7, 8])
+    transform2 = AffineTransform(A2, b2)
+
+    with pytest.raises(
+        ValueError,
+        match="Matrix dimensions of the transformations do not align for composition.",
+    ):
+        transform2.compose(transform1)
+
+
+def test_affine_transform_str():
+    A = np.array([[1, 0], [0, 1]])
+    b = np.array([1, 2])
+    transform = AffineTransform(A, b)
+    expected = "AffineTransform(A=\n[[1 0]\n [0 1]],\nb=[1 2])"
+    assert str(transform) == expected