from typing import Optional
import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.cache import KVCache


class IdentityBlock(nn.Module):
  def __call__(self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[KVCache] = None) -> mx.array:
    return x