`jax.tree.map` applies to both `Variable` and `VariableState`, despite setting `is_leaf`

[theo-brown]

Currently, applying jax.tree.map applies to both Variable and VariableState, despite setting is_leaf to select only the VariableStates.

Expected (function is not applied to Variables when selecting based on VariableState):

from flax import nnx

state = nnx.State({"param1": nnx.Param(10.0)})
print(state)
# State({
#  'param1': Param(
#    value=10.0
#  )
# })

print(isinstance(state.param1, nnx.VariableState))
# False

mapped = jax.tree.map(
    lambda x: x.value, state, is_leaf=lambda x: isinstance(x, nnx.VariableState)
)
print(mapped)
# State({
#  'param1': Param(
#    value=10.0
#  )
# })

Actual (function is applied to Variables when selecting based on VariableState)

from flax import nnx

state = nnx.State({"param1": nnx.Param(10.0)})
print(state)
# State({
#  'param1': Param(
#    value=10.0
#  )

# })
mapped = jax.tree.map(
    lambda x: x.value, state, is_leaf=lambda x: isinstance(x, nnx.VariableState)
)
print(mapped)
# State({
#  'param1': 10.0
# })

Is this behaviour intended?

Context: I found this by accident when doing the following. I need to apply a set of transforms to leaves of my State. These transforms modify/set various metadata of the leaves. Variable.replace allows you to modify metadata, however, VariableState.replace only allows modification of the value. As leaves of State can either be Variables or VariableState, I need to be able to filter between these two cases and apply different updates in each case (for Variable, I can do x.replace(transformed_value, **transformed_metadata), but for VariableState, I need to instantiate a new object, nnx.VariableState(variable_type, transformed_value, **transformed_metadata)).

[cgarciae]

Hey @theo-brown, because in your example you're using nnx.Param which a Variable, you should be checking for Variable in is_leaf:

from flax import nnx

state = nnx.State({"param1": nnx.Param(10.0)})

mapped = jax.tree.map(
    lambda x: x.value, state, is_leaf=lambda x: isinstance(x, nnx.Variable)
)
print(mapped)
# State({
#   'param1': 10.0
# })

[theo-brown]

Hi @cgarciae, thanks for your reply. I'm not sure I successfully communicated my problem. My issue is that it is incorrectly selecting Variables when I set it to only select VariableStates. I think this was a typo in my original posts MWE (I accidentally deleted the undesired 'actual' behaviour) so it's totally my fault. Let me edit and then see what you think.

For clarity, this code produces the same as your snippet, which I believe is incorrect behaviour:

from flax import nnx

state = nnx.State({"param1": nnx.Param(10.0)})

mapped = jax.tree.map(
    lambda x: x.value, state, is_leaf=lambda x: isinstance(x, nnx.VariableState)
)
print(mapped)
# State({
#   'param1': 10.0
# })

Note that I've selected based on VariableState and it's still applied the mapped function

[cgarciae]

Ah I see, I think you got a misconception about what is_leaf does, it doesn't filter-out leaves but rather it marks intermediate tree nodes as leaves, effectively stopping the traversal process and keeping them as elements that will be mapped. For your use case I think you just need to use a conditional inside your lambda if you don't want to map Variables, but this wont remove them from the tree structure.

[theo-brown]

Oh amazing, thanks so much for the clear explanation!