General question about using a state and performance

[eadadi]

Hi guys,
I want to ask a general question about the impact of using a state on performance.
---In advance - the general description is intended, sorry for that ---
I have some RL agent, that I run twice:

1. As it is (there some global state that is ignored on this run)
2. I add a function that happens on every step. this function uses the global state in the following way:
   a. The global state is a dictionary contains some values
   b. the function does very lightweight computation over these values
   c. The function returns a new state as part of its outputs (to update the global state) and preserves the functional-pure nature of the flow

The modified flow runs about 15-20% slower than the first. What might cause this? Is it simply the usage of state? The computation itself is lightweight and uses jnp. I was trying to think how the state might cause this effect, but I did not understand.

Sorry for putting things in a general manner, and I hope you might be able to give me some insights,
Thanks

[eadadi]

Thinking over it again, could it be that my function compiles over and over on every step? this will explain the performance issue

It looks like:

def func(state):
new_state = {}
NUM_A = 10
curr_a = state.get("a", jnp.zeros(NUM_A)
"""
computations with jnp
specificaly - 
next_a = ....
ret_value = ...
"""
new_state["a"] = a
return ret_value, new_state

[jakevdp]

If you change the static arguments to your function or the static attributes of the arrays passed to your function (e.g. shape, dtype), then you will trigger a re-compilation. From the code snippets you share, it's hard to tell whether this is the case.

If that doesn't answer your question, then I'd suggest putting together a minimal reproducible example showing the issue you're having – thanks!