tensorflow_backend.py 文件源码

def switch(condition, then_expression, else_expression):
    '''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.
    '''
    x_shape = copy.copy(then_expression.get_shape())
    if condition.dtype != tf.bool:
        condition = tf.cast(condition, 'bool')
    x = _cond(condition,
              lambda: then_expression,
              lambda: else_expression)
    x.set_shape(x_shape)
    return x