Keep categorical code validation inside torch.compile graphs (#1302)

> # Summary
>
> Make categorical code-range validation compatible with `torch.compile`
without registering a custom operation. Invalid codes remain checked at
runtime.
>
> # Description of the change
>
> The existing Python `if` depends on tensor values and blocks
full-graph capture. During compilation, reduce the bounds check to a
scalar, copy that scalar to CPU, and use a native PyTorch assertion. The
CPU assertion raises a recoverable error instead of a CUDA device
assertion that would poison the CUDA context.
>
> Compiled invalid input raises `RuntimeError` with the existing
message; eager execution keeps `ValueError`. Negative codes still
represent missing values. This keeps a scalar CPU
transfer/synchronization per categorical column, as the eager Python
check already synchronizes.
>
> This PR targets main independently. The [vocabulary identity
fix](https://github.com/NVIDIA/structured-data-models/tree/compile/category-vocabulary-identity)
is also needed for the tested full-graph prediction path; whole-model
compiled fit has an earlier, separate alignment blocker.
>
> <details><summary>Reproduce the compilation failure</summary>
>
> On main `b8ac2dc`, PyTorch 2.14, the compiled processor fit below
fails at the code-range Python branch. GPU validation used an NVIDIA L4,
PyTorch 2.14.0+cu130 and CUDA 13.0; CPU reproduces it with
`device="cpu"`.
>
> ```python
> import torch
> from sdm import CategoricalTensor, Stype, TableTensor
> from sdm.processing import AddCategoryCounts
>
> table = TableTensor(
>     columns={Stype.categorical: ("color",)},
>     categorical=CategoricalTensor(
> code=torch.tensor([[0], [1], [0]], dtype=torch.int32, device="cuda"),
>         categories=(torch.arange(2, device="cuda"),),
>     ),
> )
> processor = AddCategoryCounts().to("cuda")
> processor.fit(table)
>
> compiled_fit = torch.compile(processor.fit, backend="inductor",
fullgraph=True)
> compiled_fit(table)
>
> # After the fix, the same compiled graph rejects this invalid code.
> table.categorical.code[0, 0] = 2
> compiled_fit(table)
> # RuntimeError: Categorical column 'color' contains a code outside its
category vocabulary.
> ```
>
> Recorded original branch exception (`torch._dynamo.exc.Unsupported`):
>
> ```text
> Data-dependent branching
> Explanation: Detected data-dependent branching (e.g. `if
my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control
flow.
> Hint: For the common pattern `if tensor_cond: x = transform(x)` (e.g.
clamping inf/nan values), consider making the code branchless by always
applying the transform. Operations like torch.clamp, torch.nan_to_num,
and torch.where are typically no-ops on well-behaved inputs and compile
without graph breaks.
> Hint: Note: Python `or`/`and` between tensor expressions (e.g.
`tensor.any() or other_tensor.any()`) triggers implicit bool conversion.
Use `torch.logical_or`/`torch.logical_and` or the `|`/`&` operators
instead.
> Hint: This graph break is fundamental - it is unlikely that Dynamo
will ever be able to trace through your code. Consider finding a
workaround.
>   Hint: Use `torch.cond` to express dynamic control flow.
>
>   Developer debug context: attempted to jump with TensorVariable()
>
> For more details about this graph break, please visit:
https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0170.html
>
> from user code:
> File
"/Users/ardrianw/orca/workspaces/structured-data-models/kumotabular-compilation/.cache/category-branch-investigation-20261029/a/sdm/processing/base.py",
line 206, in transform
>     return self._transform(table)
>            ^^^^^^^^^^^^^^^^^^^^^^
> File
"/Users/ardrianw/orca/workspaces/structured-data-models/kumotabular-compilation/.cache/category-branch-investigation-20261029/a/sdm/processing/categorical/count.py",
line 108, in _transform
>     _check_categorical_codes(table)
> File
"/Users/ardrianw/orca/workspaces/structured-data-models/kumotabular-compilation/.cache/category-branch-investigation-20261029/a/sdm/processing/categorical/_categorical.py",
line 27, in _check_categorical_codes
>     if (codes >= category.numel()).any():
>        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>
> Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this
especially if you're reporting a bug to PyTorch). For even more
developer context, set TORCH_LOGS="+dynamo"
> ```
>
> </details>
>
> # Validation
>
> | Manual check | Result |
> | --- | --- |
> | CPU and L4 `AddCategoryCounts.fit`, `fullgraph=True` and `False` |
One Inductor graph per case; valid results match eager |
> | Invalid code after valid execution | Same compiled graph rejects the
code, without recompilation; subsequent CUDA work succeeds |
> | L4 KumoTabular predict with both fixes, `fullgraph=True` | One
graph, zero breaks, exact eager output parity |
>
> The model check used the standard small classification architecture,
default recipe, `pretrained=False`, 128 context rows and 16 query rows
after eager fit. It verifies compiler integration, not trained-model
accuracy or performance. No custom-operation registration or test
additions.

Author: aw471

sdm/processing/categorical/_categorical.py

@@ -3,6 +3,7 @@
 
 import os
 
+import torch
 from torch import Tensor
 
 from sdm import TableTensor
@@ -24,11 +25,15 @@ def _check_categorical_codes(table: TableTensor) -> None:
     categorical = table.categorical
     for index, category in enumerate(_iter_categories(categorical)):
         codes = categorical.code[..., index]
-        if (codes >= category.numel()).any():
-            raise ValueError(
-                f"Categorical column {columns[index]!r} contains a code "
-                "outside its category vocabulary."
-            )
+        message = (
+            f"Categorical column {columns[index]!r} contains a code "
+            "outside its category vocabulary."
+        )
+        if torch.compiler.is_compiling():
+            valid = (codes < category.numel()).all().to("cpu")
+            torch._assert_async(valid, message)
+        elif (codes >= category.numel()).any():
+            raise ValueError(message)
 
 
 def _check_categories(