Branch put_params_on_devices on sharding.is_fully_addressable

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -381,7 +381,12 @@ def unbox_fn(x):
       max_logging.log("🚀 Pinning all parameters to TPU HBM permanently...")
       max_logging.log("=" * 80 + "\n")
       def put_params_on_devices(params_tree, shardings_tree):
-        return jax.device_put(params_tree, shardings_tree)
+        def _put_leaf(param, sharding):
+          if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable:
+            return jax.device_put(param, sharding)
+          return device_put_replicated(param, sharding)
+
+        return jax.tree_util.tree_map(_put_leaf, params_tree, shardings_tree)
 
       max_logging.log("Putting params on TPU HBM...")
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):