code: def forward(
      model: Model[InT, InT], X: InT, is_train: bool
  ) -> Tuple[InT, Callable]:
      layer = model.layers[0]
      initial_shape = X.shape
      final_shape = list(initial_shape[:-1]) + [layer.get_dim("nO")]
      nB = X.shape[0]
      nT = X.shape[1]
      X2d = model.ops.reshape(X, (-1, X.shape[2]))
      Y2d, Y2d_backprop = layer(X2d, is_train=is_train)
      Y = model.ops.reshape3(Y2d, *final_shape)
  
      def backprop(dY: InT) -> InT:
          reshaped = model.ops.reshape2(dY, nB * nT, -1)
          return Y2d_backprop(
              model.ops.reshape3(reshaped, *initial_shape)
          )
  
      return cast(InT, Y), backprop
  
imports: from typing import Callable, List, Optional, Tuple, TypeVar, cast
  from ..config import registry
  from ..model import Model
  from ..types import Array2d, Array3d
label_desc: Use image node_id 2 for calling a global function with example usage: forward(model, X, is_train) and returns: cast,  backprop, Y2d_backprop