[Serialization] Fix: remove context manager when loading shards + handle mlx
weights
#2709
+11
−33
mlx
weights
#2709