More general remap_ids

thinc/layers/remap_ids.py

@@ -1,17 +1,22 @@
-from typing import Tuple, Callable, Sequence, Dict, Any
+from typing import Tuple, Callable, Sequence, cast
+from typing import Dict, Hashable, Union, Optional
 
 from ..model import Model
 from ..config import registry
 from ..types import Ints2d, DTypes
+from ..util import is_cupy_array, is_xp_array, to_numpy
 
-
-InT = Sequence[Any]
+InT = Union[Sequence[Hashable], Ints2d]
 OutT = Ints2d
 
 
 @registry.layers("remap_ids.v1")
 def remap_ids(
-    mapping_table: Dict[Any, int] = {}, default: int = 0, dtype: DTypes = "i"
+    mapping_table: Dict[Hashable, int] = {},
+    default: int = 0,
+    dtype: DTypes = "i",
+    *,
+    column: Optional[int] = None
 ) -> Model[InT, OutT]:
     """Remap string or integer inputs using a mapping table, usually as a
     preprocess before embeddings. The mapping table can be passed in on input,
@@ -21,7 +26,12 @@ def remap_ids(
     return Model(
         "remap_ids",
         forward,
-        attrs={"mapping_table": mapping_table, "dtype": dtype, "default": default},
+        attrs={
+            "mapping_table": mapping_table,
+            "dtype": dtype,
+            "default": default,
+            "column": column
+        },
     )
 
 
@@ -31,11 +41,22 @@ def forward(
     table = model.attrs["mapping_table"]
     default = model.attrs["default"]
     dtype = model.attrs["dtype"]
+    column = model.attrs["column"]
+    if column is not None:
+        inputs = cast(Ints2d, inputs)
+        inputs = inputs[:, column]
+    # elements of cupy arrays are 0-dimensional arrays
+    # not the integers stored in the original mapper.
+    if is_cupy_array(inputs):  # type: ignore
+        inputs = to_numpy(inputs)  # type: ignore
     values = [table.get(x, default) for x in inputs]
     arr = model.ops.asarray2i(values, dtype=dtype)
     output = model.ops.reshape2i(arr, -1, 1)
 
     def backprop(dY: OutT) -> InT:
-        return []
+        if is_xp_array(inputs):
+            return model.ops.xp.empty(dY.shape)  # type: ignore
+        else:
+            return []
 
     return output, backprop

thinc/tests/layers/test_layers_api.py

@@ -126,7 +126,8 @@ def assert_data_match(Y, out_data):
     # ("CauchySimilarity.v1", {}, (array2d, array2d), array1d),
     ("ParametricAttention.v1", {}, ragged, ragged),
     ("SparseLinear.v1", {}, (numpy.asarray([1, 2, 3], dtype="uint64"), array1d, numpy.asarray([1, 1], dtype="i")), array2d),
-    ("remap_ids.v1", {"dtype": "f"}, ["a", 1, 5.0], array2dint)
+    ("remap_ids.v1", {"dtype": "f"}, ["a", 1, 5.0], array2dint),
+    ("remap_ids.v1", {"dtype": "f", "column": 1}, numpy.array([[1, 2, 3], [4, 5, 6]]).T, array2dint)
     # fmt: on
 ]