How to train a TensorFlow model having branches that need to train alternately according to the data type؟

Question

Suppose I have two data a and b, and I want to compute the similarity between them using the dot product. However, there are two types of a, a1 and a2; such that, the feature of a1 and a2 are obtained by module_a1 and module_a2 respectively. The feature of b is obtained from module_b.

My question is how can I implement OR operation that existed in the attached figure. Please, see this image to clarify the idea.

It means if the type of a is a1, use module_a1 to generate features of a in the training (i.e. use it in forward and backward propagation), else if the type is a2, use module_a2. Therefore, in some iterations, module_a1 gets trained, and in the other iterations, module_a2 gets trained.

So, could you help me, with how can I implement this idea in TensorFlow? Thanks