jax lax.while_loop and lax.fori_loop with multiple conditions and statements

[PoulomiPradhan]

I have python for loop and while loop with multiple conditions and multiple statements inside the loop. I would like to convert it to lax.while_loop and lax.fori_loop. So would it contain multiple body funcs and cond funcs ?
Can someone please help how to do that ?

[jakevdp]

Hi - thanks for the question! It's hard to answer such things in abstract: could you edit your question to add a small minimal reproducible example showing the situation you have in mind?

[PoulomiPradhan]

Thanks for your answer. This time I figured it out by myself. I just created another function instead of multiple statements and used that in jax.scan()