TensorFlow : Transform Tensor based on conditions of features

Thank you very much for reading my question. I'm quite new to TensorFlow so sorry if my problem doesn't completely make sense, I have this regression problem :

A = Input(shape=(8))
A0 = A[:,0:4]
A1 = A[:,4:8]
A0 = layers.Dense(12)(A0)
A1 = layers.Dense(12)(A1)
Z = layers.Concatenate()([A0,A1])
Z = layers.Dense(1)(Z)
model = Model(inputs=A, outputs=Z)

wherein the input there are two sets of features, observation A0 && A1 ie.[temperature, humidity, UV, pollutants] from two separate devices.

1. if I knew there is an interdependency between these two sets and I want to transform A0 based on the output of Dense()(A1) something like ----> A0 depends on the features found within [A1 ?!@?#-> A0]
2. and I also want final layer Z to only depends on A0

what kind of method should I use? would something like this make sense? or Tf.Cond or if conditions?

A0 = layers.Dense(12)(A0)
A1 = layers.Dense(12)(A1)
Z = layers.Concatenate()([A0,A1])
Z = layers.Dense(12)(Z)
Z = A0 + Z
Z = layers.Dense(1)(Z)

I'm looking for something that seems better looking and elegant, or may you please point out a path like some relevant studies that I can look at?

I think my main problem is I don't even know what to look at, as it is not exactly an if condition/tf.case problem

Thank you very much for your time

Many Thanks