Repository: qthequartermasterman/hypothesis-torch

Issue:
Feature Request: Add the ability to generate named tensors to `tensor_strategy`
https://pytorch.org/docs/stable/named_tensor.html

Discussion:

Patch:
diff --git a/hypothesis_torch/tensor.py b/hypothesis_torch/tensor.py
index dd3415e..825a352 100644
--- a/hypothesis_torch/tensor.py
+++ b/hypothesis_torch/tensor.py
@@ -51,6 +51,7 @@ def tensor_strategy(
     pin_memory: bool | st.SearchStrategy[bool] | None = None,
     layout: torch.layout | st.SearchStrategy[torch.layout] | None = None,
     memory_format: torch.memory_format | st.SearchStrategy[torch.memory_format] | None = None,
+    names: bool | st.SearchStrategy[bool] = False,
 ) -> torch.Tensor:
     """A strategy for generating PyTorch tensors.
 
@@ -76,6 +77,8 @@ def tensor_strategy(
         memory_format: The memory format of the tensor. If None, a suitable default will be inferred based on the other
             arguments. Note that channel_last memory formats are only supported for 4D tensors and channel_last_3d
             memory formats are only supported for 5D tensors.
+        names: Whether to give explicit names to the tensor's dimensions, using the "Named Tensors" API. `names` will
+            default to False until the Named Tensors API is no longer experimental.
 
     Returns:
         A strategy for generating PyTorch tensors.
@@ -183,6 +186,22 @@ def tensor_strategy(
     # Pyright/mypy falsely reports an error here on py3.9 torch 2.1.2.
     tensor = tensor.to(memory_format=memory_format)  # type: ignore
 
+    if isinstance(names, st.SearchStrategy):
+        names = draw(names)
+    if names:
+        # Named tensors only supported on CPU/CUDA.
+        hypothesis.assume(tensor.device.type in {"cpu", "cuda"})
+        # Named tensors only support strided tensors
+        hypothesis.assume(layout == torch.strided)
+        num_dimensions = tensor.ndim
+        individual_dimensions_should_have_name = draw(st.tuples(*[st.booleans() for _ in range(num_dimensions)]))
+        # At least one dimension needs to be named, but not all of them have to be.
+        hypothesis.assume(any(individual_dimensions_should_have_name))
+        individual_dimensions_names = [
+            f"dim{i}" if should_name else None for i, should_name in enumerate(individual_dimensions_should_have_name)
+        ]
+        tensor = tensor.rename(*individual_dimensions_names)
+
     return tensor
 
 


Test patch:
diff --git a/tests/unit/test_tensor.py b/tests/unit/test_tensor.py
index 6e22627..22440d5 100644
--- a/tests/unit/test_tensor.py
+++ b/tests/unit/test_tensor.py
@@ -405,3 +405,40 @@ class TestTensor(unittest.TestCase):
         )
 
         assert tensor.dim_order() == tuple(range(tensor.dim()))
+
+    @hypothesis.given(st.data())
+    def test_names_false_has_no_named_dimensions(self, data: st.DataObject) -> None:
+        """Test that if the names parameter is False, no dimensions have names."""
+        tensor = data.draw(
+            hypothesis_torch.tensor_strategy(
+                dtype=torch.float32,
+                shape=(1, 2, 3),
+                names=False,
+            )
+        )
+
+        assert all(name is None for name in tensor.names)
+
+    @hypothesis.given(st.data())
+    def test_names_true_has_named_dimension(self, data: st.DataObject) -> None:
+        """Test that if the names parameter is True, at least one dimension has a name."""
+        tensor = data.draw(
+            hypothesis_torch.tensor_strategy(
+                dtype=torch.float32,
+                shape=(1, 2, 3),
+                names=True,
+            )
+        )
+
+        assert any(name is not None for name in tensor.names)
+
+    @hypothesis.given(st.data())
+    def test_names_strategy_is_valid(self, data: st.DataObject) -> None:
+        """Test that if the names parameter is a strategy, only valid combinations are found."""
+        _ = data.draw(
+            hypothesis_torch.tensor_strategy(
+                dtype=torch.float32,
+                shape=(1, 2, 3),
+                names=st.booleans(),
+            )
+        )