code: def forward(
      model: Model[InT, OutT], X: InT, is_train: bool
  ) -> Tuple[OutT, Callable]:
      column: int = model.attrs["column"]
      layer = model.layers[0]
      if X.size < 2:
          return layer(X, is_train)
      keys = X[:, column]
      if not isinstance(keys, numpy.ndarray):
          keys = keys.get()  # pragma: no cover
      uniq_keys, ind, inv, counts = layer.ops.xp.unique(
          keys,
          return_index=True,
          return_inverse=True,
          return_counts=True,
      )
      counts = model.ops.reshape2i(counts, -1, 1)
      X_uniq = X[ind]
      Y_uniq, bp_Y_uniq = layer(X_uniq, is_train)
      Y = Y_uniq[inv].reshape((X.shape[0],) + Y_uniq.shape[1:])
      uniq_shape = tuple(Y_uniq.shape)
  
      def backprop(dY: OutT) -> InT:
          dY_uniq = layer.ops.alloc2f(*uniq_shape)
          layer.ops.scatter_add(dY_uniq, layer.ops.asarray_i(inv), dY)
          d_uniques = bp_Y_uniq(dY_uniq)
          # This confusing bit of indexing "ununiques"
          return (d_uniques / counts)[inv]
  
      return Y, backprop
  
imports: from typing import Callable, Optional, Tuple
  import numpy
  from ..config import registry
  from ..model import Model
  from ..types import Floats2d, Ints2d
label_desc: Use image node_id 2 for calling a global function with example usage: forward(model, X, is_train) and returns: Y,  backprop, layer, unknown