535d15402492fc10579a3211dc83c74f2741cc31,my/tensorflow/general.py,,flatten,#Any#Any#,121

Before Change

    if keep == 0:
        out_shape = [-1]
    else:
        out_shape = [-1] + tensor.get_shape().as_list()[-keep:]
    flat = tf.reshape(tensor, out_shape)
    return flat

After Change

def flatten(tensor, keep):
    fixed_shape = tensor.get_shape().as_list()
    start = len(fixed_shape) - keep
    left = reduce(mul, [fixed_shape[i] or tf.shape(tensor)[i] for i in range(start)])
    out_shape = [left] + [fixed_shape[i] or tf.shape(tensor)[i] for i in range(start, len(fixed_shape))]
    flat = tf.reshape(tensor, out_shape)
    return flat

Instances