Reject negative PyTorch one-hot class counts

src/pyrecest/backend_support/_pytorch_one_hot_scalar_contract.py

@@ -6,6 +6,7 @@
 
 
 _INTEGER_MESSAGE = "num_classes must be an integer"
+_NONNEGATIVE_MESSAGE = "num_classes must be non-negative"
 
 
 def _is_boolean_take_axis(axis, torch_module) -> bool:
@@ -20,17 +21,20 @@ def _is_boolean_take_axis(axis, torch_module) -> bool:
 
 
 def _normalize_num_classes(num_classes, torch_module) -> int:
-    """Return a non-boolean integer ``num_classes`` value."""
+    """Return a non-negative, non-boolean integer ``num_classes`` value."""
     if isinstance(num_classes, bool) or type(num_classes).__name__ == "bool_":
         raise TypeError(f"{_INTEGER_MESSAGE}, not boolean")
     if torch_module.is_tensor(num_classes):
         if num_classes.ndim != 0 or num_classes.dtype == torch_module.bool:
             raise TypeError(_INTEGER_MESSAGE)
         num_classes = num_classes.item()
     try:
-        return _operator_index(num_classes)
+        normalized = _operator_index(num_classes)
     except TypeError as exc:
         raise TypeError(_INTEGER_MESSAGE) from exc
+    if normalized < 0:
+        raise ValueError(_NONNEGATIVE_MESSAGE)
+    return normalized
 
 
 def _patch_pytorch_take_axis_contract(pytorch_backend, torch_module) -> None:

tests/backend_support/test_pytorch_one_hot_scalar_contract.py

@@ -38,6 +38,14 @@ def _scalar_contract_code(target_module):
     else:
         raise AssertionError("boolean num_classes was accepted")
 
+for bad_num_classes in (-1, torch.tensor(-1)):
+    try:
+        target.one_hot([0], bad_num_classes)
+    except ValueError as exc:
+        assert "num_classes must be non-negative" in str(exc)
+    else:
+        raise AssertionError("negative num_classes was accepted")
+
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("meta")
 device_result = target.one_hot(torch.tensor(2, device=device), 4)
 assert device_result.device.type == device.type