Key Error Issue that I don't fully understand

[eadadi]

Hello,
I am relatively new to Jax, and there is something I need to do.

I want to understand why the following code yields Key Error ('A').:

@jax.jit                                                                                                                                                                                                       
def r(state):
  def fa(state):
    return 0
  def fb(state):
    return state['A']
  return lax.cond('A' not in state, fa, fb, state)
state = {}
state['A'] = r(state)          

The compiled version should fit any version of dictionary, isn't it?

The reason I need something like this to work is because I want to get something like:

@jax.jit                                                                                                                                                                                                       
def r(state):
  def fa(state):
    return 0
  def fb(state):
    return state['A']
  return lax.cond('A' not in state, fa, fb, state)
state = {}
state['A'] = r(state)
r(state) #now the condition takes the other branch, and it should be valid because 'A' is in state      

[jakevdp]

Hi - thanks for the question! The issue is that static operations like dict key lookups happen at trace-time, and lax.cond will always trace both branches (though at runtime it will execute only one branch). You can fix this by changing fb to account for the fact that it will be traced regardless of the condition; e.g.

  def fb(state):
    return state.get('A', 0)

But backing up a step, given that your condition is a statically-known quantity, there's no reason to use lax.cond and to push the logic into the compiler. You could simplify your function by evaluating this static condition at trace time at the top level; for example:

@jax.jit                                                                                                                                                                                                       
def r(state):
  return state.get('A', 0)

[eadadi]

Thank you very much Jake, your help is very appreciated.

So you say that both branches are traced, which makes sense, but with what value? the fallback one?
I'll give you example. I have something like this:

@jax.jit
  def r(arg):
    cond = arg.get('STATE', None)
    def fa(state):
      return 0
    def fb(state):
      return 1 + state.get('x', None)
return lax.cond(cond is not None, fa, fb, cond)                                                                                                                                                            
                                                                                                                                                                                                           
arg = {'STATE': {'x': 8}}
print(r(arg))

And I get error that NoneType has no attribute 'get', which makes me thinks that tracing is with the fall back value,
because if I do

cond = arg.get('STATE',{})

it works.

BUT PLEASE NOTE: the above example does not yields that error and it actually works, I can't really tell what is the different from the project I am working on, but maybe you'll still be able to help me understand what's happenning

[jakevdp]

So you say that both branches are traced, which makes sense, but with what value?

It's traced with an abstract value representing the shape and dtype of the value that will be encountered at runtime.

There's some more background on tracing and how to think about JAX's execution/transformation model at How to think in JAX.

In your followup example, cond is not None is known statically at trace-time, so there's no reason to pass it to lax.cond. I'd suggest implementing that function like this instead:

@jax.jit
def r(arg):
  if 'STATE' in arg:
    return 1 + arg['STATE']['x']
  else:
    return 0

[eadadi]

I started with something similar to your approach, but then I had a phenomenon where it seemed to me like the same random number was always sampled (inside a logic in one of the branches).
I guessed that some value remained fixed during runtime, so I tried to write it the way I posted here.

I will update soon.

(p.s. Yes, I've read this article. I still not full grasp it unfurtunately)

[jakevdp]

It sounds like you're calling numpy.random somewhere in your code. This will be executed only once and stored as a compile-time constant. If you want runtime random values in JAX, you need to use jax.random rather than numpy.random.

[eadadi]

I will check it out again.
I'm resolving this since you answered my question.
Thanks for your help