Merge remote-tracking branch 'origin/main' into spec-with-neg-shape

# Conflicts: # test/test_specs.py
pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit caa5f4178d4203771dba2fc45248c56618303034
diff --git a/test/test_specs.py b/test/test_specs.py
@@ -3483,6 +3483,14 @@ def test_all(self):
         assert spec.is_in(data)
 
 
+class TestNonTensorSpec:
+    def test_sample(self):
+        nts = NonTensorSpec(shape=(3, 4))
+        assert nts.one((2,)).shape == (2, 3, 4)
+        assert nts.rand((2,)).shape == (2, 3, 4)
+        assert nts.zero((2,)).shape == (2, 3, 4)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)