fix (stateful)(utilities): adding FlaxPreTrainedModel as one of the…

… bases. This is because `FlaxPreTrainedModel` is a standalone class that doesn't directly inherit from a native Flax module.

ivy/stateful/utilities.py

@@ -23,7 +23,7 @@ def _is_submodule(obj, kw):
             "keras.src.engine.base_layer.Layer",
             "keras.src.layers.layer.Layer",
         ),
-        "flax": ("flax.nnx.nnx.module.Module",),
+        "flax": ("flax.nnx.nnx.module.Module","transformers.modeling_flax_utils.FlaxPreTrainedModel"),
     }[kw]
     try:
         for bc in type(obj).mro():
@@ -730,7 +730,7 @@ def forward(self, x):
     def _compute_module_dict_jax(model, prefix=""):
         _module_dict = dict()
         for key, value in model.__dict__.items():
-            if isinstance(value, nnx.Module):
+            if isinstance(value, nnx.Module) and value != model:
                 if not hasattr(value, "named_parameters"):
                     _module_dict.update(
                         _compute_module_dict_jax(value, prefix=f"{key}.")