Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_factory

Author: kylesayrs

src/compressed_tensors/transform/utils/hadamard.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import math
-from typing import Tuple
+from typing import Optional, Tuple
 
 import numpy
 import torch
@@ -58,21 +58,20 @@ def deterministic_hadamard_matrix(size: int) -> torch.Tensor:
 # https://github.com/Dao-AILab/fast-hadamard-transform/tree/master
 
 
-def random_hadamard_matrix(size: int) -> torch.Tensor:
+def random_hadamard_matrix(
+    size: int, gen: Optional[torch.Generator] = None
+) -> torch.Tensor:
     """
     Produces a randomly generated Hadamard matrix.
     See https://cornell-relaxml.github.io/quip-sharp/ ,
     Section "Randomized Hadamard Transformation"
 
-    :param size: The dimension of the matrix. Matrix generated will have dimensions
-        (size, size)
-
+    :param size: The dimension of the hamadard matrix
+    :param gen: Optional generator random values
+    :return: randomly generated hadamard matrix
     """
-    # TODO: potentially update to add "seed" as an arugment, to allow
-    # the matrix generated to be reproducible
-
     # Benefits: support other shapes / non powers of 2, support randomization
-    Q = torch.randint(low=0, high=2, size=(size,)).to(torch.float64)
+    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=torch.float64)
     Q = Q * 2 - 1
     Q = torch.diag(Q)
     return _matmul_hadU(Q) / math.sqrt(size)

tests/test_transform/utils/test_hadamard.py

@@ -49,6 +49,30 @@ def test_random_hadamard_matrix_compliant(size):
     assert torch.equal(product, torch.eye(size))
 
 
+def test_random_hadamard_generator():
+    generator = torch.Generator().manual_seed(42)
+    one = random_hadamard_matrix(2048, generator)
+    two = random_hadamard_matrix(2048, generator)
+
+    one_true = torch.tensor(
+        [
+            [-1, -1, -1],
+            [+1, -1, +1],
+            [-1, -1, +1],
+        ]
+    )
+    two_true = torch.tensor(
+        [
+            [-1, -1, -1],
+            [-1, +1, -1],
+            [+1, +1, -1],
+        ]
+    )
+
+    assert torch.all(one[:3, :3].sign() == one_true.sign())
+    assert torch.all(two[:3, :3].sign() == two_true.sign())
+
+
 @pytest.mark.parametrize(
     "size",
     [1024],