How do dim and nclasses interact? If dim > 1, is nclasses vector-valued / broadcastable? Or do we require that all indices have the same max index? I'm OK with either approach (vector-valued nclasses is more representationally flexible, but fixed nclasses supports conversion to one-hot tensor format), but we should document it clearly here.