Use native jax.device_put matching reference FLUX/SDXL models

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -381,12 +381,7 @@ 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):
-        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)
+        return jax.device_put(params_tree, shardings_tree)
 
       max_logging.log("Putting params on TPU HBM...")
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):