Fix AQT conv general construction.

PiperOrigin-RevId: 641070533

chirp/models/efficientnet.py

@@ -23,7 +23,7 @@
 from typing import Callable, NamedTuple
 
 from aqt.jax.v2 import aqt_conv_general
-from aqt.jax.v2 import config as aqt_cfg
+from aqt.jax.v2 import aqt_dot_general
 from chirp.models import layers
 from flax import linen as nn
 import flax.typing as flax_typing
@@ -133,7 +133,7 @@ class OpSet:
         head_activation=nn.hard_swish,
         dot_general=jax.lax.dot_general,
         conv_general_dilated=aqt_conv_general.make_conv_general_dilated(
-            aqt_cfg.DotGeneralRaw.make_conv_general_dilated()
+            aqt_conv_general.conv_general_dilated_make(spatial_dimensions=2)
         ),
     ),
 }