Added util to print pytree structure and sharding rule based on calla…

…ble policy.

scalax/sharding.py

@@ -103,6 +103,25 @@ def get_partition_spec(name, leaf):
         return named_tree_map(get_partition_spec, pytree, sep='/')
 
 
+class PolicyShardingRule(ShardingRule):
+    """ Create PartitionSpec for a pytree with a callable policy. """
+
+    def __init__(self, policy):
+        """ Create a PolicyShardingRule with a callable policy.
+
+        Args:
+            policy: A callable that takes a tree path and a leaf tensor as input
+                and returns a PartitionSpec.
+        """
+        self.policy = policy
+
+    def apply(self, pytree):
+        """ Returns a pytree of PartitionSpec according to the policy. """
+        def get_partition_spec(name, leaf):
+            return self.policy(name, leaf)
+        return named_tree_map(get_partition_spec, pytree, sep='/')
+
+
 class MeshShardingHelper(object):
     """ Helper class for creating jit sharding jax functions with sharding rules. """
     global_mesh_helper = None

scalax/utils.py

@@ -106,3 +106,10 @@ def named_tree_map(f, tree, *rest, is_leaf=None, sep=None):
         is_leaf=is_leaf
     )
 
+
+def print_pytree_structure(tree, sep='/', is_leaf=None):
+    def print_fn(path, val):
+        shape = f'shape: {val.shape if hasattr(val, "shape") else "none"}'
+        dtype = f'dtype: {val.dtype if hasattr(val, "dtype") else "none"}'
+        print(f'{path}: {shape}, {dtype}')
+    named_tree_map(print_fn, tree, is_leaf=is_leaf, sep=sep)