def train_step(
    model: Any,
    state: TrainState,
    batch: Dict[str, jnp.ndarray],
    rng: np.ndarray,
    learning_rate_fn: Callable[[int], float],
    weight_decay: float,
    grad_clip_max_norm: Optional[float] = None