Variable vs VariableState

[carlesoctav]

I have a question regarding the difference between the two. This question came to mind because of the nnx.get_partition_spec() functionality.

def get_partition_spec(tree: A) -> A:
    """Extracts a PartitionSpec tree from a PyTree containing ``Variable`` values."""

    def _maybe_replicate(x):
        if hasattr(x, 'shape'):
            return PartitionSpec()
        else:
            return None

    def f(x):
        if isinstance(x, (variablelib.VariableState, variablelib.Variable)):
            if hasattr(x, 'sharding') and x.sharding:
                if core_spmd.get_logical_axis_rules() or hasattr(x, 'sharding_rules'):
                    context_rules = core_spmd.get_logical_axis_rules()
                    local_rules = getattr(x, 'sharding_rules', ())
                    rules = core_spmd.composite_rules(context_rules, local_rules)
                    return x.replace(
                        PartitionSpec(*core_spmd.from_sharding_rules(x.sharding, rules))
                    )
                return x.replace(PartitionSpec(*x.sharding))
            else:
                return x.replace(_maybe_replicate(x.value))

        return _maybe_replicate(x)

    return jax.tree.map(
        f, tree, is_leaf=lambda x: isinstance(x, variablelib.VariableState)
    )

What I’m trying to do is create function a that automatically create PartitionSpec by inferring the shape of the parameters and the mesh configuration.

When I try to do this:

state = nnx.state(model)
jax.tree.map(
    lambda x: x.shape, is_leaf=lambda x: isinstance(x, variablelib.VariableState)
)

It results in an error: x has no attribute shape. But when I run the following code:

state = nnx.state(model)
jax.tree.map(
    lambda x: x.shape, is_leaf=lambda x: isinstance(x, variablelib.Variable)
)

It runs fine.

If I understand correctly, the variablelib.Variable corresponds to param.value, so it has the .shape and even .sharding properties (even though it's still the unsharded one). But what exactly is VariableState used for, and what is the corresponding part in the nnx.Module class? Also, why does nnx.get_partition_spec() perform the tree mapping with VariableState as the definition of leaf?

[carlesoctav]

i understand, now.