filter_shape: Sequence[int],
      is_qat_model: bool = False,
  ) -> module.Module:
    class SimpleGatherAndConvModel(module.Module):
      """A simple model with a single gather and a conv2d."""

      def __init__(self):
        """Initializes a SimpleGatherAndConvModel."""
        self.embedding_w = np.random.randn(1024, 3, 4, 3).astype('f4')