Does passing big pytrees into each function slow down Jax?

[MRiabov]

Sup Jax,

I'm building a huge numerical simulator program. Dozens of files, often 500+ lines of code of pure logic, and increase even further.

I'm mostly translating code from another big C++ program with pure OOP, and they have the following approach:

Have a central large object that stores almost everything about the program - the state of the program, most of it nested,
And this central large object is passed to almost every high-level method in the code, which then modifies it in-place.

While obviously the "modifies and returns" part does not work for jax, wouldn't passing big amounts of data (multiple megabytes of state) to every function slow Jax down? (maybe it's how jit works, idk.)

I would want to avoid it, but I sometimes end up passing 7-8 inputs into the function that can all be replaced with a single pytree.

Cheers.