jax.lax.cond and the None datatype?

[rog77]

Hi, I wonder if anyone can advise on this please?

I run the following code:

import jax

A = 2
B = None
C = None


def branching_fn1(A, B = None, C = None):

  def positive_branch1(A):
    return A

  def negative_branch1(A):

    def branching_fn2(A, B, C):
      return jax.lax.cond((C is None), positive_branch2, negative_branch2, A)

    def positive_branch2(A):
      A = A + B # rewrite this in jax primitives
      return A  

    def negative_branch2(A):
      A = A + (C*B) # rewrite this in jax primitives
      return A  


    return branching_fn2(A, B, C)  
    
  return jax.lax.cond((B is None), positive_branch1, negative_branch1, A)

output = branching_fn1(A, B, C)

print(output)

And get variations on the error: unsupported operand type(s) for +: 'DynamicJaxprTracer' and 'NoneType' - regardless of if the conditional path should be taken, given the starting data.

I note that if I run the code below it works as I would expect. Since I'm trying to convert numpy code to jax, and it already uses "None", understanding how to resolve this would be really helpful. B & C in the original code may be None or vector/scalar - In the working code it is possible that I could remove the first conditional, and have two versions of jit'd functions dispatched with a pythoninc "if", but that doesn't help with the second one. And help/advice/links would me most appreciated. I am still very new to Jax, so please don't be afraid to explain it to me like I'm a small child. Thanks in advance!

import jax

A = 2
B = -1
C = -1


def branching_fn1(A, B = -1, C = -1):

  def positive_branch1(A):
    return A

  def negative_branch1(A):

    def branching_fn2(A, B, C):
      return jax.lax.cond((C == -1), positive_branch2, negative_branch2, A)

    def positive_branch2(A):
      A = A + B # rewrite this in jax primitives
      return A  

    def negative_branch2(A):
      A = A + (C*B) # rewrite this in jax primitives
      return A  


    return branching_fn2(A, B, C)  
    
  return jax.lax.cond((B == -1), positive_branch1, negative_branch1, A)

output = branching_fn1(A, B, C)

print(output) 

[jakevdp]

The problem is that lax.cond is not a short-circuiting operation: in order to push the computation down to XLA, it traces both the true_fun and false_fun to construct an intermediate representation, even if the final code will not need to evaluate one of the branches.

But even so, I have trouble following the logic of the first code-block: for example, it seems that if C is not None and A = None, then even with short-circuiting logic positive_branch2 would attempt to execute A = A + B, which would be ill-defined even without respect to JAX.

Another thing you could consider: if you're branching on python identity comparisons (as opposed to value comparisons), instead of using lax.cond you could use a simple if/else construction. For example, this should be fine for negative_branch, even with JIT/autodiff:

  def negative_branch1(A):
    if C is None:
      return A + B
    else:
      return A + C * B

[rog77]

Thanks for the response, below is the code snippet I am trying to adapt. I am quite keen to see it going in XLA so will ask about using -1 rather than None. I appreciate what you say about it being ill defined as I wrote it. Was I correct in defining the second branch function inside of the first negative branch? I'm happy to mark off the answer you already gave, unless you have any further advice.

While I am on the subject, A will be a result of a prior calculation - would I need to use blockuntilready on that, or does jax know not to attempt to run a calculation on A until a prior step is complete, but will evaluate the if block and enqueue the next calculation? This is for my understanding, I'm wary about using the pythonic "if", if it will slow the code down as it will run a lot of times, but if this is evaluated whilst the previous function is running on the device, I suppose it it is less of a problem.

Thanks again!

A will always be a vector/device array.
B may be a vector or None.
C may be a (positive) scalar or None.

if B is not None:
  if C is not None:
    A += C*B
  else:
    A += B

[jakevdp]

That seems fine as written:

import jax

@jax.jit
def f(A, B, C):
  if B is not None:
    if C is not None:
      A +=  C*B
    else:
      A += B
  return A

A = 1.0
B = None
C = 2.0

print(f(1.0, None, None))
# 1.0
print(f(1.0, 2.0, None))
# 3.0
print(f(1.0, 2.0, 3.0))
# 7.0

You can't condition if statements on the values in arrays within JIT, but conditioning on Python object identity is fine.

[rog77]

Thanks :-)