[BUG] Problems with multidimensional nvec in MultiCategorical #2486

Activity

1. matteobettini

   added

   bugSomething isn't workingSomething isn't working

   on Oct 10, 2024
2. matteobettini

   assigned

   vmoens

   on Oct 10, 2024
3. kostasrigatos commented on Oct 3, 2026

   

   kostasrigatos

   on Oct 3, 2026

   More actions

   Context

   This builds on the recent work in #4493 / #4504.

   MultiCategorical.to_one_hot_spec() is fixed and it now correctly handles 2D or higher nvec configurations by unbinding and stacking specs:

   a = MultiCategorical([[3, 2], [1, 4]])
   a.to_one_hot_spec()  # succeeds, returns a Stacked/MultiOneHot spec

   This fix is unrelated to #4493/#4504 (those PRs only touch to_one_hot, not to_one_hot_spec), so it looks like this part of the issue was resolved independently at some point and the issue was never updated to reflect it.

   The Problem

   MultiCategorical.to_one_hot() — the value-encoding method, not the spec-conversion one above — is still broken for ndim > 1 inputs. It throws a TypeError because it loops through the nvec as if it were always 1D:

   a = MultiCategorical([[3, 2], [1, 4]])
   a.to_one_hot(a.rand())
   # TypeError: one_hot(): argument 'num_classes' (position 2) must be int, not Tensor

   This is a distinct code path from what #4504 addresses — that PR's to_one_hot fix is scoped to batched/singleton specs, and explicitly notes that the heterogeneous-row case (this issue) is "a different case" and is left untouched. This PR covers that remaining gap.

   Proposed Fix

   Essentially a mirror of the logic to_one_hot_spec() already uses:

   • When self.ndim > 1, unbind the spec and the value along the appropriate axis.
   • Recursively apply to_one_hot to each resulting pair.
   • If all resulting one-hot tensors have the same width, stack them into a dense tensor.
   • If widths differ (genuinely heterogeneous nvec, e.g. [[1,2],[3,4]]), raise a ValueError — a dense tensor can't represent rows of different width, which is also why to_one_hot_spec() falls back to a lazy Stacked wrapper instead of a flat MultiOneHot in that case.

   This also covers the uniform-batched case from #4493/#4504 (e.g. MultiCategorical([3,2], shape=(4,2))) for free, since unbind doesn't distinguish uniform from heterogeneous rows — both go through the same path.

   Testing

   Added parametrized tests covering: heterogeneous widths (expects ValueError), matching widths (expects successful dense output), 3-level nested nvec, and an extra leading batch dimension beyond self.shape (to catch axis-misalignment bugs).

   Full suite: pytest test/test_specs.py — 1816 passed, 7 skipped (CUDA-only). ufmt/flake8 clean.

   Future Scope (left out of this fix)

   1. MultiCategorical.to_categorical is currently a no-op (return val) rather than a real one-hot → categorical decoder — so there's no full round-trip for the ndim > 1 case.
   2. Stacked has no to_categorical/to_categorical_spec override, so the lazy spec returned by to_one_hot_spec() for heterogeneous rows can't itself decode a sample back to categorical form.

   PR incoming with the fix + tests. @vmoens @Nicholas022400701 — happy to adjust the approach if either of you sees it differently.
4. bsprenger

   closed this as completedon Oct 8, 2026