@@ -0,0 +1,1302 @@
+from typing import Union, Sequence, Tuple, TypeVar, cast, Callable, overload%0Afrom ..types import ArrayXd, FloatsXd, IntsXd%0Afrom ..model import Model%0A%0A%0AAxisIndex = Union%5Bint, slice, Sequence%5Bint%5D%5D%0AIndex = Union%5BAxisIndex, Tuple%5BAxisIndex, ...%5D%5D%0A%0A%0Adef array_getitem(index: Index) -%3E Model%5BArrayXd, ArrayXd%5D:%0A    %22%22%22Index into input arrays, and return the subarrays.%0A    %0A    The %60index%60 object can %0A    %22%22%22%0A    return Model(%22array-getitem%22, forward, attrs=%7B%22index%22: index%7D)%0A%0A%0Adef floats_getitem(index: Index) -%3E Model%5BFloatsXd, FloatsXd%5D:%0A    %22%22%22Index into input arrays, and return the subarrays.%0A    %0A    This delegates to %60array_getitem%60, but allows type declarations.%0A    %22%22%22%0A    return Model(%22floats-getitem%22, forward, attrs=%7B%22index%22: index%7D)%0A%0A%0Adef ints_getitem(index: Index) -%3E Model%5BIntsXd, IntsXd%5D:%0A    %22%22%22Index into input arrays, and return the subarrays.%0A    %0A    This delegates to %60array_getitem%60, but allows type declarations.%0A    %22%22%22%0A    return Model(%22ints-getitem%22, forward, attrs=%7B%22index%22: index%7D)%0A%0A%0Adef forward(model, X, is_train):%0A    index = model.attrs%5B%22index%22%5D%0A    shape = X.shape%0A    dtype = X.dtype%0A    %0A    def backprop_get_column(dY):%0A        dX = model.ops.alloc(*shape, dtype=dtype) %0A        dX%5Bindex%5D = dY %0A        return dX%0A%0A    Y = X%5Bindex%5D %0A    return Y, backprop_get_column%0A

Add array getitem