flax switch error

[jakubMitura14]

Hello I have couple big branches that I earlier tried to switch between using jax switch. switch function worked but some modules did not with error message that I should not mix jax and flax transformations - hence I am trying to use flax version

first I prepare all possible versions with different indexes

import functools
from flax import linen as nn

def main_fun_op(self,curried,index):
    pass

main_fun_ops=list(map(lambda index: functools.partial(self.main_fun_op,index=index) ,range(self.r_x_all.shape[0])))

next I invoke and it DO work

main_fun_ops[0](curried) 

switch do not work

nn.switch(indexx, self.main_fun_ops, self, curried)

error

main_fun_op() got multiple values for argument 'index'

I can not make sense of it as the index is already set by partial application

main_fun_op - is complex
All functions are executed from within nn.Module
main_fun_ops - is prepared in setup function of a module switch is in call

Thanks for help !!

[jakevdp]

It's unclear from your partial code, but it appears that main_func_op is a method of a class (because you're referring to it as self.main_func_op) in which case the self argument is implicit so you probably need something like:

nn.switch(indexx, self.main_fun_ops, curried)

This would be easier to answer if you could provide a minimal reproducible example. Also, since since this is a flax-specific question, you'll probably get better answers if you ask at http://github.com/google/flax/

[jakubMitura14]

Thanks Yes it is all about self argument; when I do it as below it works

        def a0(self,curried):
            return self.main_fun_op(curried,0)
        def a1(self,curried):
            return self.main_fun_op(curried,1)
        def a2(self,curried):
            return self.main_fun_op(curried,2)
        def a3(self,curried):
            return self.main_fun_op(curried,3)
        def a4(self,curried):
            return self.main_fun_op(curried,4)
        def a5(self,curried):
            return self.main_fun_op(curried,5)
       
        # return nn.switch(index,self.main_fun_ops,self,curried)
        # main_fun_ops[0]()
        return nn.switch(indexx, [ a0,a1,a2,a3,a4,a5 ], self, curried)

Thanks !!