How are variable size pytrees jitted when variable sized arrays need to be a static argument?

[aeftimia]

Like the question says, Jax can jit functions of variable-sized pytree arguments and intermediate values but not variable-sized ndarrays. What's going on under the hood that explains that behavior?

[jakevdp]

Can you edit your question with a brief example of what you're asking about? I'm having trouble understanding from the short description what exactly you are asking.

[aeftimia]

Hey I tried to create some test cases and it looks like I misunderstood. It seems you can pass variable size pytrees, and variable sized arrays as arguments to a function to be jitted. But you cannot create a pytree or array of variable size within the function body. So the two scenarios appear consistent.

[jakevdp]

Great, glad you were able to answer your own question 😀