why are we able to use 'if' inside a 'pmap' train_step?

[yiyixuxu]

I saw this pattern in the imagenet example - can anyone help me understand why are we able to use if inside a pmap here? the argument state here is a flax TrainState type

I really appreciate your help!

YiYi

@jax.pmap
def train_step(state,...):
   if state.dynamic_scale:
       ....
 
train_step(state,...)

https://github.com/google/flax/blob/main/examples/imagenet/train.py#L128

[andsteing]

You can use Python control logic inside a jit/pmap, as long as that control flow is not value dependent.

In this case, the if statement is about the value of state.dynamic_scale which can be None, if no dynamic scaling is used, or am instance of dynamic_scale_lib.DynamicScale, if dynamic scaling is used. In either case, the value will be set from the config (see lines 227-230 in code snippet below), and is not changed during the training. In other words, jit/pmap will produce a single compiled version of the code.

flax/examples/imagenet/train.py

Lines 222 to 244 in 6fd2446

	def create_train_state(rng, config: ml_collections.ConfigDict,
	model, image_size, learning_rate_fn):
	"""Create initial training state."""
	dynamic_scale = None
	platform = jax.local_devices()[0].platform
	if config.half_precision and platform == 'gpu':
	dynamic_scale = dynamic_scale_lib.DynamicScale()
	else:
	dynamic_scale = None
	params, batch_stats = initialized(rng, image_size, model)
	tx = optax.sgd(
	learning_rate=learning_rate_fn,
	momentum=config.momentum,
	nesterov=True,
	)
	state = TrainState.create(
	apply_fn=model.apply,
	params=params,
	tx=tx,
	batch_stats=batch_stats,
	dynamic_scale=dynamic_scale)
	return state

[yiyixuxu]

ohhh thank you! i got it now. it is not actually a bool type