Support transformer_engine v2.16.

PiperOrigin-RevId: 936695878

Author: Google-ML-Automation

src/maxtext/layers/quantizations.py

@@ -1012,19 +1012,27 @@ def _wrap(self, f, name=None):
     """
 
     import transformer_engine.jax  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+    from transformer_engine.common import recipe  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
 
     fp8_recipe = self._recipe
 
     class TEWrapper(transformer_engine.jax.flax.module.TransformerEngineBase):
       """Wrapper module for TransformerEngine quantization."""
 
-      def generate_quantizer_set(self, postfix: str = ""):
+      def generate_quantizer_set(self, postfix: str = "",
+          variable_collection: str | None = None,
+          quantization_checkpoint_name: str | None = None,
+          fp8_recipe: recipe.Recipe | None = None,
+          n_groups: int | None = None,
+      ):
+
         OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient"
         return super().generate_quantizer_set(  # pytype: disable=wrong-keyword-args
             postfix=postfix,
             variable_collection=OVERWRITE_WITH_GRADIENT,
             quantization_checkpoint_name="quantization",
             fp8_recipe=fp8_recipe,
+            n_groups=n_groups,
         )
 
       @nn.compact