From 24b0bc832c6805b5bf93d96f95531673105454ac Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 15:33:35 -0700
Subject: [PATCH 1/2] Record the nonempty bound of categorical unique outputs

---
 sdm/tensor/categorical.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py
index e9466fee7..aea660ffb 100644
--- a/sdm/tensor/categorical.py
+++ b/sdm/tensor/categorical.py
@@ -329,6 +329,10 @@ def from_tensor(cls, tensor: Tensor) -> Self:
                 for column in tensor.unbind(dim=-1)
             ]
         )
+        # A nonempty column always has at least one unique value.
+        if tensor.numel() > 0:
+            for category in categories:
+                torch._check(category.numel() > 0)
         return cls(
             code=torch.stack(values, dim=-1),
             categories=categories,

From 567e527334bfc513227cab640c66d73f1822eab8 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 19:23:15 -0700
Subject: [PATCH 2/2] Clarify why categorical vocabulary bounds aid compilation

---
 sdm/tensor/categorical.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py
index aea660ffb..2d2172762 100644
--- a/sdm/tensor/categorical.py
+++ b/sdm/tensor/categorical.py
@@ -329,7 +329,7 @@ def from_tensor(cls, tensor: Tensor) -> Self:
                 for column in tensor.unbind(dim=-1)
             ]
         )
-        # A nonempty column always has at least one unique value.
+        # Tell torch.compile that nonempty columns have at least one category.
         if tensor.numel() > 0:
             for category in categories:
                 torch._check(category.numel() > 0)