Using nnx.call with multiple states

[renecotyfanboy]

Is there any syntax to use nnx.call with multiple states? I don't see ways to do :

nnx.call((graph_def, state_1, state_2)) 

in the docs, and it doesn't work anyway. I am currently doing the following within a jit compiled function:

graphdef, parameters, tables = nnx.split(self, nnx.Param, ...)

parameters = set_parameters(parameters) # Edit the param state

result = nnx.merge(graphdef, parameters, tables).method() # Maybe nnx.call().method() could be better here

Can I stay with this syntax, is there any overhead in using merge instead of call?
Thank you in advance for your help !

[vfdev-5]

@renecotyfanboy it is a bit unclear the purpose of

nnx.call((graph_def, state_1, state_2)) 

with multiple states.

Split and merge paradigm is correct in your example. nnx.call behind the scene is doing that:

flax/flax/nnx/graph.py

Lines 2861 to 2867 in f73aea5

	def pure_caller(accessor: DelayedAccessor, *args, **kwargs):
	node = merge(*graphdef_state)
	method = accessor(node)
	out = method(*args, **kwargs)
	return out, split(node)
	return CallableProxy(pure_caller) # type: ignore

[cgarciae]

@renecotyfanboy nnx.call does not work for multiple states as its supposed to be used with the outputs of nnx.split(model). You can consider using nnx.merge_state to create a single state for call.