About jnp.where to control two different parts of network

[lucasliunju]

Hi, I would like to use jnp.where to control two parts of network. More specially, I use:

g = jnp.where(t<30000, jax.tree_map(lambda x: jax.lax.pmean(x, axis_name='batch'), g), g)

To contro whether we need to communicate with other hosts at each step "t".

I try to print "t":
Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=0/1)>

However, the error is:
ValueError: entry not a 2- or 3- tuple

Could you give me some advice, thank you!

Yong

[jheek]

I'm a bit confused by the error message but your problem is that you cannot apply jnp.where on a tree so should probably use:

g = jax.tree_map(lambda x: jnp.where(t<30000, lax.pmean(x, 'batch'), g), g)