from typing import Dict, Callable, Tuple, Optional, Any, Union, cast

from ..model import Model
from ..config import registry
from ..types import Ints2d, Floats2d, Ints1d
from ..initializers import uniform_init
from ..util import get_width, partial


InT = Union[Ints1d, Ints2d]
OutT = Floats2d


@registry.layers("Embed.v1")
def Embed(
    nO: Optional[int] = None,
    nV: Optional[int] = None,
    *,
    column: int = 0,
    initializer: Callable = uniform_init,
    dropout: Optional[float] = None
) -> Model[InT, OutT]:
    """Map integers to vectors, using a fixed-size lookup table."""
    attrs: Dict[str, Any] = {"column": column}
    if dropout is not None:
        attrs["dropout_rate"] = dropout
    return Model(
        "embed",
        forward,
        init=partial(init, initializer),
        attrs=attrs,
        dims={"nO": nO, "nV": nV},
        params={"E": None},
    )


def forward(model: Model[InT, OutT], ids: InT, is_train: bool) -> Tuple[OutT, Callable]:
    nV = model.get_dim("nV")
    vectors = cast(Floats2d, model.get_param("E"))
    column: int = model.attrs["column"]
    dropout = model.attrs.get("dropout_rate")
    input_shape = tuple(ids.shape)
    if ids.ndim == 2:
        ids1d = ids[:, column]  # type: ignore
    else:
        ids1d = cast(Ints1d, ids)
    ids1d *= ids1d < nV
    output = vectors[ids1d.astype("i")]
    drop_mask = cast(Floats2d, model.ops.get_dropout_mask((vectors.shape[1],), dropout))
    output *= drop_mask

    def backprop(d_output: OutT) -> InT:
        d_output *= drop_mask
        d_vectors = model.ops.alloc2f(*vectors.shape)
        model.ops.scatter_add(d_vectors, ids1d, d_output)
        model.inc_grad("E", d_vectors)
        dX: InT = model.ops.alloc(input_shape, dtype=ids1d.dtype)
        return dX

    return output, backprop


def init(
    initializer: Callable,
    model: Model[InT, OutT],
    X: Optional[InT] = None,
    Y: Optional[OutT] = None,
) -> Model[InT, OutT]:
    if Y is not None:
        model.set_dim("nO", get_width(Y))
    shape = (model.get_dim("nV"), model.get_dim("nO"))
    model.set_param("E", initializer(model.ops, shape))
    return model

@@ -63,16 +63,81 @@
 , cast%0A%0A
+from .chain import chain%0Afrom .array_getitem import ints_getitem%0A
 from ..m
@@ -208,17 +208,35 @@
 ort Ints
-2
+1d, Ints2d, Floats1
 d, Float
@@ -337,29 +337,14 @@
 T = 
-Union%5BInts1d, Ints2d%5D
+Ints1d
 %0AOut
@@ -481,15 +481,28 @@
 mn: 
+Optional%5B
 int
+%5D
  = 
-0
+None
 ,%0A  
@@ -693,32 +693,36 @@
 tr, 
-Any%5D = %7B%22column%22: column
+Union%5BNone, int, float%5D%5D = %7B
 %7D%0A  
@@ -795,21 +795,37 @@
     
-return Model(
+model = Model( # type: ignore
 %0A   
@@ -986,16 +986,169 @@
 ,%0A    )%0A
+    model = model if column is None else chain(ints_getitem(column), model)%0A    model.attrs%5B%22column%22%5D = column%0A    return cast(Model%5BInT, OutT%5D, model)%0A%0A
 %0A%0Adef fo
@@ -1234,37 +1234,8 @@
 e%5D:%0A
-    nV = model.get_dim(%22nV%22)%0A
     
@@ -1289,261 +1289,114 @@
     
-column: int = model.attrs%5B%22column%22%5D%0A    dropout = model.attrs.get(%22dropout_rate%22)%0A    input_shape = tuple(ids.shape)%0A    if ids.ndim == 2:%0A        ids1d = ids%5B:, column%5D  # type: ignore%0A    else:%0A        ids1d = cast(Ints1d, ids)%0A    ids1d *= ids1d %3C nV
+nO = vectors.shape%5B1%5D%0A    nN = ids.shape%5B0%5D%0A    dropout: Optional%5Bfloat%5D = model.attrs.get(%22dropout_rate%22)
 %0A   
@@ -1420,22 +1420,8 @@
 %5Bids
-1d.astype(%22i%22)
 %5D%0A  
@@ -1437,33 +1437,33 @@
 sk = cast(Floats
-2
+1
 d, model.ops.get
@@ -1477,32 +1477,18 @@
 t_mask((
-vectors.shape%5B1%5D
+nO
 ,), drop
@@ -1687,18 +1687,16 @@
 ors, ids
-1d
 , d_outp
@@ -1748,21 +1748,16 @@
       dX
-: InT
  = model
@@ -1770,39 +1770,13 @@
 lloc
-(input_shape, dtype=ids1d.dtype
+1i(nN
 )%0A  
@@ -2066,16 +2066,18 @@
 im(%22nV%22)
++1
 , model.

Update embed layer