fix (stateful)(utilities): added a guard to handle cases wherein `jax…

….Array's` might be present inside the `.v` dict. We also call `setattr` to ensure the update is reflected on the translated model as well.

ivy/stateful/utilities.py

@@ -201,10 +201,16 @@ def _maybe_update_flax_layer_weights(
                 )
                 params2[name] = getattr(layer, flax_name)
                 continue
-
-            params2[name].value = jnp.asarray(
-                params1_np, dtype=params2[name].value.dtype
-            )
+
+            if isinstance(params2[name], nnx.Variable):
+                params2[name].value = jnp.asarray(
+                    params1_np, dtype=params2[name].value.dtype
+                )
+            else:
+                params2[name] = jnp.asarray(
+                    params1_np, dtype=params2[name].dtype
+                )
+                setattr(model2, name, params2[name])
 
         for name in buffers1:
             layer, weight_name = _retrive_layer(model2, key_mapping[name])
@@ -250,6 +256,7 @@ def _maybe_update_flax_layer_weights(
 
             else:
                 buffers2[name] = jnp.asarray(buffers1_np, dtype=buffers2[name].dtype)
+                setattr(model2, name, buffers2[name])
 
     # Check if the parameters and buffers are the same
     for name in params1: