optax.multi_transform + nnx.State/nnx.Optimizer troubles

[yklcs]

optax.multi_transform defines multiple transforms with a Mapping[Hashable, GradientTransformation] and uses a PyTree or function to map parameters to the key.

Using optax.multi_transform with nnx.Optimizer means said mapping of type nnx.State is needed.
nnx.State is typed to use StateLeaf which means we can't use string or integer keys.
While ignoring typing does work, it feels brittle and might end up broken later.
Is there any other solution for this problem?

tx = optax.multi_transform(
    {
        "weights": optax.adamw(learning_rate, momentum),
        "biases": optax.adamw(learning_rate, momentum),
        },
        # this doesn't work:
        # {
        #     "weights": "weights",
        #     "biases": "biases",
        # },
        # this does, but is it safe?:
        nnx.State({
            "weights": "weights",
            "biases": "biases"
        })
    }
)

[cgarciae]

Interesting. You are right that State is currently restrictive regarding the value type. I think it can be safely generalized to be generic over the value type e.g. State[V] and then have V be StateLeaf for current APIs graph APIs but users can just change the type for these kind of situations.

To answer your question I think its safe to do what you are showing. I'll going to turn this into an issue.

[yklcs]

Closing in favor of #3955