remove "return state" for functions which will be executed in vmap; recover randkey as args in mutation methods
we recommend to use output_transform; change hyperparameters (strong) in XOR example;