fixing lint

Author: parambole

MaxText/layers/models.py

@@ -21,7 +21,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax.ad_checkpoint import checkpoint_name
 from jax.sharding import Mesh
 
 from flax import linen as nn
@@ -31,12 +30,10 @@
 from MaxText import max_logging
 from MaxText.inference import page_manager
 from MaxText.layers import linears
-from MaxText.layers import quantizations
 from MaxText.layers import pipeline
 from MaxText import maxtext_utils
 from MaxText import multimodal_utils
 from MaxText.layers.blocks import DecoderLayer
-from MaxText.layers.attentions import Attention
 from MaxText.layers.normalizations import RMSNorm
 from MaxText.layers.embeddings import PositionalEmbedding, Embed
 from MaxText.layers.quantizations import AqtQuantization as Quant