def cond_shift_and_scale(self, inputs, cond_inputs, is_cond, layer):
    if not is_cond:
      return inputs
    cond_out = layer(cond_inputs)
    if self.cond_scale:
      scale, shift = tf.split(cond_out, num_or_size_splits=2, axis=-1)
      scale = self.cond_act_func(scale)
      shift = self.cond_act_func(shift)
      inputs *= scale
      inputs += shift
    else:
      inputs += cond_out
    return inputs