Fix scalar CUDA hashing and enable it for CUDA tensor dumps

tests/test_cuda_hashing_contracts.py

@@ -0,0 +1,192 @@
+"""Backend routing and numeric parity for the existing tensor-hash algorithm."""
+
+import unittest
+from collections.abc import Sequence
+from unittest.mock import patch
+
+import torch
+
+from traincheck.instrumentor import dumper
+from traincheck.instrumentor.proxy_wrapper import hash as tensor_hashing
+
+
+def reference_hash(values: Sequence[int]) -> int:
+    """Independent Python-integer recurrence with explicit signed-64 overflow."""
+    accumulator = 0
+    for value in values:
+        accumulator = (accumulator * 6364136223846793005 + value + 1) % (2**64)
+    return accumulator - 2**64 if accumulator >= 2**63 else accumulator
+
+
+def first_quantized_row(tensor: torch.Tensor) -> list[int]:
+    """Preserve the historical quantization and first-row selection contract."""
+    cpu = tensor.detach().cpu()
+    if cpu.is_floating_point():
+        cpu = (cpu * 1e8).to(torch.int64)
+    else:
+        cpu = cpu.to(torch.int64)
+    if cpu.ndim == 0:
+        return [int(cpu.item())]
+    if cpu.ndim == 1:
+        return [int(value) for value in cpu.tolist()]
+    return [int(value) for value in cpu.reshape(-1, cpu.shape[-1])[0].tolist()]
+
+
+class CudaTaggedTensor(torch.Tensor):
+    """Exercise routing without claiming that this CPU-backed fixture is a GPU."""
+
+    is_cuda: bool = True
+
+
+class TensorDumpRoutingTests(unittest.TestCase):
+    def setUp(self) -> None:
+        selected = patch.dict(
+            dumper.tensor_dump_format,
+            dump_tensor_hash=True,
+            dump_tensor_stats=False,
+            dump_tensor_full=False,
+        )
+        selected.start()
+        self.addCleanup(selected.stop)
+
+    def test_cuda_tag_routes_to_cuda_without_cpu_fallback(self) -> None:
+        value = torch.zeros(1).as_subclass(CudaTaggedTensor)
+        with patch.object(dumper, "tensor_hash", return_value=37) as calculate:
+            self.assertEqual(dumper.dump_tensor(value), 37)
+        calculate.assert_called_once_with(value, with_parallel=True, with_cuda=True)
+
+    def test_non_cuda_tensor_keeps_cpu_hash_route(self) -> None:
+        value = torch.zeros(1)
+        with patch.object(dumper, "tensor_hash", return_value=41) as calculate:
+            self.assertEqual(dumper.dump_tensor(value), 41)
+        calculate.assert_called_once_with(value, with_parallel=True, with_cuda=False)
+
+    def test_cuda_hash_failure_propagates_without_silent_fallback(self) -> None:
+        value = torch.zeros(1).as_subclass(CudaTaggedTensor)
+        error = RuntimeError("CUDA hash failed")
+        with patch.object(dumper, "tensor_hash", side_effect=error) as calculate:
+            with self.assertRaises(RuntimeError) as caught:
+                dumper.dump_tensor(value)
+        self.assertIs(caught.exception, error)
+        calculate.assert_called_once_with(value, with_parallel=True, with_cuda=True)
+
+    def test_other_dump_modes_and_non_tensors_keep_their_behavior(self) -> None:
+        value = torch.tensor([1.0, 2.0])
+        with (
+            patch.dict(
+                dumper.tensor_dump_format, dump_tensor_hash=False, dump_tensor_full=True
+            ),
+            patch.object(dumper, "tensor_hash") as calculate,
+        ):
+            self.assertEqual(dumper.dump_tensor(value), [1.0, 2.0])
+            self.assertIsNone(dumper.dump_tensor(object()))
+            calculate.assert_not_called()
+        cached = object()
+        with (
+            patch.dict(dumper.tensor_dump_format, dump_tensor_stats=True),
+            patch.object(dumper, "tensor_stats", return_value=cached),
+            patch.object(dumper, "tensor_hash") as calculate,
+        ):
+            self.assertIs(dumper.dump_tensor(value), cached)
+            calculate.assert_not_called()
+
+    def test_cpu_hash_matches_independent_recurrence_without_mutation(self) -> None:
+        for shape in ((), (4,), (3, 4), (2, 3, 4)):
+            elements = 1
+            for dimension in shape:
+                elements *= dimension
+            value = (torch.arange(elements, dtype=torch.int64) - 3).reshape(shape)
+            original = value.clone()
+            rng = torch.get_rng_state().clone()
+            with self.subTest(shape=shape):
+                self.assertEqual(
+                    tensor_hashing.tensor_hash(value, with_cuda=False),
+                    reference_hash(first_quantized_row(value)),
+                )
+                self.assertTrue(torch.equal(value, original))
+                self.assertTrue(torch.equal(torch.get_rng_state(), rng))
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "actual CUDA device required")
+class CudaHashValueTests(unittest.TestCase):
+    def check_value(self, value: torch.Tensor) -> None:
+        self.assertTrue(value.is_cuda)
+        original = value.clone()
+        cpu_rng = torch.get_rng_state().clone()
+        cuda_rng = torch.cuda.get_rng_state(value.device).clone()
+        expected = reference_hash(first_quantized_row(value))
+        cpu_hash = tensor_hashing.tensor_hash(value, with_cuda=False)
+        cuda_hash = tensor_hashing.tensor_hash(value, with_cuda=True)
+        self.assertEqual(cuda_hash, expected)
+        self.assertEqual(cuda_hash, cpu_hash)
+        self.assertTrue(torch.equal(value, original))
+        self.assertTrue(torch.equal(torch.get_rng_state(), cpu_rng))
+        self.assertTrue(torch.equal(torch.cuda.get_rng_state(value.device), cuda_rng))
+
+    def test_cuda_scalar_hashes_match_cpu_for_supported_dtypes(self) -> None:
+        for dtype in (
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+            torch.float64,
+            torch.int8,
+            torch.int16,
+            torch.int32,
+            torch.int64,
+            torch.uint8,
+        ):
+            for scalar in (0, 1, -1) if dtype != torch.uint8 else (0, 1, 255):
+                # Small representable half values avoid quantization to infinity.
+                value = scalar / 16384 if dtype == torch.float16 else scalar
+                with self.subTest(dtype=dtype, scalar=value):
+                    self.check_value(torch.tensor(value, dtype=dtype, device="cuda"))
+
+    def test_cuda_vector_matrix_and_3d_match_the_same_recurrence(self) -> None:
+        for shape in ((4,), (3, 4), (2, 3, 4)):
+            elements = 1
+            for dimension in shape:
+                elements *= dimension
+            for dtype in (torch.int64, torch.float32, torch.float64, torch.bfloat16):
+                with self.subTest(shape=shape, dtype=dtype):
+                    value = (torch.arange(elements, device="cuda") - 3).to(dtype)
+                    self.check_value(value.reshape(shape))
+
+    def test_cuda_noncontiguous_tensor_and_empty_vector_keep_existing_results(
+        self,
+    ) -> None:
+        value = torch.arange(24, dtype=torch.int64, device="cuda").reshape(4, 6)
+        self.check_value(value[:, ::2])
+        self.check_value(value.t())
+        self.check_value(torch.empty(0, dtype=torch.float32, device="cuda"))
+
+    def test_real_cuda_dumper_uses_cuda_hash_and_preserves_gradients(self) -> None:
+        value = torch.tensor(-0.5, device="cuda", requires_grad=True)
+        value.grad = torch.tensor(2.0, device="cuda")
+        original_gradient = value.grad.clone()
+        with (
+            patch.dict(
+                dumper.tensor_dump_format,
+                dump_tensor_hash=True,
+                dump_tensor_stats=False,
+                dump_tensor_full=False,
+            ),
+            patch.object(
+                tensor_hashing,
+                "hash_tensor_cuda",
+                wraps=tensor_hashing.hash_tensor_cuda,
+            ) as cuda_backend,
+            patch.object(
+                tensor_hashing,
+                "hash_tensor_cpu",
+                side_effect=AssertionError("CPU fallback"),
+            ),
+        ):
+            self.assertEqual(
+                dumper.dump_tensor(value), reference_hash(first_quantized_row(value))
+            )
+            cuda_backend.assert_called_once()
+        self.assertTrue(torch.equal(value.grad, original_gradient))
+
+
+if __name__ == "__main__":
+    unittest.main()

traincheck/instrumentor/dumper.py

@@ -315,17 +315,20 @@ def tensor_stats(tensor: torch.Tensor):
     return result
 
 
-def dump_tensor(value):
+def dump_tensor(value: object) -> object:
+    """Encode a tensor using the selected dump mode; non-tensors return None.
+
+    The legacy statistics path can return application-cached metadata, which
+    has not been runtime-validated here, so the result contract remains object.
+    CUDA hashing preserves the same recurrence as the CPU path while avoiding
+    its full tensor transfer. CPU and MPS tensors retain CPU hashing.
+    """
     param_list = None
     if isinstance(value, torch.Tensor):
         if tensor_dump_format["dump_tensor_stats"]:
             param_list = tensor_stats(value)
         elif tensor_dump_format["dump_tensor_hash"]:
-            if value.is_cuda:
-                param_list = tensor_hash(value, with_parallel=True, with_cuda=False)
-            else:
-                # TODO: support quick hashing methods for MPS tensors
-                param_list = tensor_hash(value, with_parallel=True, with_cuda=False)
+            param_list = tensor_hash(value, with_parallel=True, with_cuda=value.is_cuda)
 
         elif tensor_dump_format["dump_tensor_full"]:
             param_list = value.detach().flatten().tolist()

traincheck/instrumentor/proxy_wrapper/hash.py

@@ -22,9 +22,17 @@ def cuda_hash_kernel(data, hash_values, multiplier, increment):
             hash_values[idx] = hash_value
 
 
-def hash_tensor_cuda(x):
+def hash_tensor_cuda(x: Tensor) -> int:
+    """Apply the existing first-row hash to an integer CUDA tensor.
+
+    Scalars represent a one-element row, matching the CPU implementation.
+    Empty vectors retain their existing zero hash; zero-row matrices remain
+    outside this function's supported input shape contract.
+    """
     # if x is more than 2D, flatten it to 2D
-    if x.ndim > 2:
+    if x.ndim == 0:
+        x = x.reshape(1, 1)
+    elif x.ndim > 2:
         x = x.flatten(start_dim=0, end_dim=-2)
     elif x.ndim == 1:
         # if x is 1D, add a dimension to make it 2D (n x 1)
@@ -40,8 +48,8 @@ def hash_tensor_cuda(x):
         x, hash_values, MULTIPLIER, INCREMENT
     )
 
-    x = hash_values.copy_to_host()
-    return int(x[0])
+    host_hash_values = hash_values.copy_to_host()
+    return int(host_hash_values[0])
 
 
 def hash_tensor_cpu(x: torch.Tensor) -> int:
@@ -95,7 +103,7 @@ def _reduce_last_axis(x: Tensor) -> Tensor:
 def tensor_hash(x: Tensor, with_parallel: bool = True, with_cuda: bool = True) -> int:
     if hasattr(x, "_traincheck_tensor_hash"):
         return x._traincheck_tensor_hash
-    if hasattr(x, "to_local"):  
+    if hasattr(x, "to_local"):
         x = x.to_local()
     if with_parallel:
         if x.dtype in [