Repository: explosion/thinc

Issue:
to_categorical discards input shape
While porting a PyTorch example to thinc, I noticed that `to_categorical` does not produce the same results as TensorFlow or PyTorch. Specifically, thinc discards the shape when dealing with more than 2d inputs.

While looking at other libraries to compare results, I noticed there are differences between how the output shape is calculated based on whether you're using TF/Torch helpers, or `keras.utils.to_categorical` 😅

Here's a snippet demonstrating the differences:
```python
>>> import numpy as np
>>> data = np.random.randint(0,12, size=(2,3,1))

>>> from thinc.api import to_categorical
>>> to_categorical(data, n_classes=12).shape
(6, 12)

>>> import torch
>>> torch.nn.functional.one_hot(torch.from_numpy(data),12).shape
torch.Size([2, 3, 1, 12])

>>> import tensorflow as tf
>>> tf.one_hot(data, depth=12).shape
TensorShape([2, 3, 1, 12])

>>> import keras
>>> keras.utils.to_categorical(data, 12).shape
(2, 3, 12)
```

@honnibal what should be the correct behavior for thinc? Perhaps `[2, 3, 1, 12]` since TensorFlow and PyTorch agree on that output?


Discussion:

Resolving patch:
diff --git a/thinc/util.py b/thinc/util.py
--- a/thinc/util.py
+++ b/thinc/util.py
@@ -51,7 +51,7 @@
 except ImportError:  # pragma: no cover
     has_mxnet = False
 
-from .types import ArrayXd, ArgsKwargs, Ragged, Padded, Floats2d, IntsXd
+from .types import ArrayXd, ArgsKwargs, Ragged, Padded, FloatsXd, IntsXd
 
 
 def get_array_module(arr):  # pragma: no cover
@@ -193,18 +193,20 @@ def copy_array(dst: ArrayXd, src: ArrayXd) -> None:  # pragma: no cover
         numpy.copyto(dst, src)
 
 
-def to_categorical(Y: IntsXd, n_classes: Optional[int] = None) -> Floats2d:
+def to_categorical(Y: IntsXd, n_classes: Optional[int] = None) -> FloatsXd:
     # From keras
     xp = get_array_module(Y)
     if xp is cupy:  # pragma: no cover
         Y = Y.get()
+    keep_shapes: List[int] = list(Y.shape)
     Y = numpy.array(Y, dtype="int").ravel()
-    if not n_classes:
-        n_classes = numpy.max(Y) + 1
+    if n_classes is None:
+        n_classes = int(numpy.max(Y) + 1)
+    keep_shapes.append(n_classes)
     n = Y.shape[0]
     categorical = numpy.zeros((n, n_classes), dtype="float32")
     categorical[numpy.arange(n), Y] = 1
-    return xp.asarray(categorical)
+    return xp.asarray(categorical).reshape(keep_shapes)
 
 
 def get_width(


Test patch:
diff --git a/thinc/tests/layers/test_mxnet_wrapper.py b/thinc/tests/layers/test_mxnet_wrapper.py
--- a/thinc/tests/layers/test_mxnet_wrapper.py
+++ b/thinc/tests/layers/test_mxnet_wrapper.py
@@ -40,7 +40,10 @@ def X(input_size: int) -> Array2d:
 @pytest.fixture
 def Y(answer: int, n_classes: int) -> Array2d:
     ops: Ops = get_current_ops()
-    return to_categorical(cast(IntsXd, ops.asarray([answer])), n_classes=n_classes)
+    return cast(
+        Array2d,
+        to_categorical(cast(IntsXd, ops.asarray([answer])), n_classes=n_classes),
+    )
 
 
 @pytest.fixture
diff --git a/thinc/tests/test_util.py b/thinc/tests/test_util.py
--- a/thinc/tests/test_util.py
+++ b/thinc/tests/test_util.py
@@ -43,7 +43,6 @@ def test_array_module_cpu_gpu_helpers():
     assert not is_numpy_array((1, 2))
 
 
-@pytest.mark.xfail(reason="¯/_(ツ)_/¯")
 def test_to_categorical():
     # Test without n_classes
     one_hot = to_categorical(numpy.asarray([1, 2], dtype="i"))
@@ -52,7 +51,7 @@ def test_to_categorical():
     # https://github.com/keras-team/keras/blob/master/tests/keras/utils/np_utils_test.py
     nc = 5
     shapes = [(1,), (3,), (4, 3), (5, 4, 3), (3, 1), (3, 2, 1)]
-    expected_shapes = [(1, nc), (3, nc), (4, 3, nc), (5, 4, 3, nc), (3, nc), (3, 2, nc)]
+    expected_shapes = [(1, nc), (3, nc), (4, 3, nc), (5, 4, 3, nc), (3, 1, nc), (3, 2, 1, nc)]
     labels = [numpy.random.randint(0, nc, shape) for shape in shapes]
     one_hots = [to_categorical(label, nc) for label in labels]
     for label, one_hot, expected_shape in zip(labels, one_hots, expected_shapes):