How should I achieve "if" statement in pmap?

[Cccjl219]

I use the pamp in my code, but I need to use a judgment statement to achieve the comparsion of two numbers. But the error is jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/2)>

I know it is related to jit complie because I use the pmap function. I try a lot way, such as static argnums, control flow, and so on, but I still failed.
I hope that I can obtain the help. Thank you!!

[Cccjl219]

I revise my code to avoid "if" using lax.cond and jnp. where. But I meet new error~crying.

def compare_loss(loss, max_loss):
    return jnp.where((loss-max_loss)>0,1,0)
is_update_grad = compare_loss(loss, max_loss)
max_loss, max_params = jax.lax.cond(is_update_grad, true_fun1(loss, noised_model, max_loss, max_params),false_fun1(max_loss, max_params))

  File "C:\Users\HP\Desktop\Update GNP\gnp\training\training_core.py", line 176, in get_gnp_gradient
    jax.lax.cond(is_update_grad, true_fun1(loss, noised_model, max_loss, max_params),false_fun1(max_loss, max_params))
TypeError: unhashable type: 'DynamicJaxprTracer'

Can you help me? Thanks

[jakevdp]

Hi - thanks for the question!

There are two options that may work for you depending on the context:

1. Mark the quantity as static: this would make sense if you are doing something like if flag where flag is a boolean value passed directly to the function
2. Use lax.cond or jnp.where: this would make sense if you are conditioning on a quantity that is computed from dynamic input, such as values within an array.

If you need more specific advice, I'd suggest sharing a minimal reproducible example of the code that leads to the error you're seeing.