code: def switch(condition, then_expression, else_expression):
      
      x_shape = copy.copy(then_expression.get_shape())
      x = tf.cond(tf.cast(condition, 'bool'),
                  lambda: then_expression,
                  lambda: else_expression)
      x.set_shape(x_shape)
      return x
docs: Switches between two operations depending on a scalar value (int or bool).
      Note that both `then_expression` and `else_expression`
      should be symbolic tensors of the *same shape*.
  
      # Arguments
          condition: scalar tensor.
          then_expression: TensorFlow operation.
          else_expression: TensorFlow operation.