from keras.models import Model
from keras.layers import Input, Convolution2D, Activation, Flatten, Dense, Reshape, Lambda
from keras.layers.advanced_activations import PReLU
from keras.engine.topology import merge
from keras.optimizers import Nadam
import keras.backend as K

from eva.layers.residual_block import ResidualBlockList
from eva.layers.masked_convolution2d import MaskedConvolution2D

def PixelCNN(input_shape, filters, blocks, build=True):
    width, height, channels = input_shape

    input_map = Input(shape=input_shape, name='input_map')

    model = MaskedConvolution2D(filters, 7, 7, mask='A', border_mode='same')(input_map)
    model = PReLU()(model)

    model = ResidualBlockList(model, filters, blocks)

    model = MaskedConvolution2D(filters, 1, 1)(model)
    model = PReLU()(model)

    model = MaskedConvolution2D(3*256, 1, 1)(model)
    model = PReLU()(model)

    # TODO: Make it scalable to any amount of channels.

    red = MaskedConvolution2D(256, 1, 1)(model)
    red = Reshape((input_shape[0] * input_shape[1], 256))(red)
    red = Activation('softmax', name='red')(red)

    green = MaskedConvolution2D(256, 1, 1)(model)
    green = Reshape((input_shape[0] * input_shape[1], 256))(green)
    green = Activation('softmax', name='green')(green)

    blue = MaskedConvolution2D(256, 1, 1)(model)
    blue = Reshape((input_shape[0] * input_shape[1], 256))(blue)
    blue = Activation('softmax', name='blue')(blue)

    # TODO: Make is scalable to any amount of channels.

    if build:
        model = Model(input=input_map, output=[red, green, blue])
        model.compile(optimizer=Nadam(clipnorm=1., clipvalue=1.),
                      loss={    'red':   'sparse_categorical_crossentropy',
                                'green': 'sparse_categorical_crossentropy',
                                'blue':  'sparse_categorical_crossentropy'})

    return model

@@ -262,16 +262,35 @@
 .backend
+.tensorflow_backend
  as K%0A%0Af
@@ -407,16 +407,69 @@
 lution2D
+%0Afrom eva.layers.fuckingsoftmax import FuckingSoftmax
 %0A%0Adef Pi
@@ -1019,43 +1019,187 @@
     
-red = MaskedConvolution2D(256, 1, 1
+model = Reshape((input_shape%5B0%5D, input_shape%5B1%5D, 256, input_shape%5B2%5D))(model)%0A%0A    # TODO: Make it scalable to any amount of channels.%0A%0A    red = Lambda(lambda x: x%5B:, :, :, :, 0%5D
 )(mo
@@ -1332,37 +1332,41 @@
 n = 
-MaskedConvolution2D(256, 1
+Lambda(lambda x: x%5B:, :, :, :
 , 1
+%5D
 )(mo
@@ -1508,37 +1508,41 @@
 e = 
-MaskedConvolution2D(256, 1, 1
+Lambda(lambda x: x%5B:, :, :, :, 2%5D
 )(mo
@@ -1843,33 +1843,8 @@
 dam(
-clipnorm=1., clipvalue=1.
 ),%0A 
@@ -1883,22 +1883,20 @@
 red':   
-'spars
+imag
 e_catego
@@ -1905,33 +1905,32 @@
 cal_crossentropy
-'
 ,%0A              
@@ -1956,22 +1956,20 @@
 green': 
-'spars
+imag
 e_catego
@@ -1986,17 +1986,16 @@
 sentropy
-'
 ,%0A      
@@ -2033,58 +2033,792 @@
 ':  
-'sparse_categorical_crossentropy'%7D)%0A%0A    return model
+image_categorical_crossentropy%7D)%0A%0A    return model%0A%0Afrom keras.backend.common import _EPSILON%0Afrom keras.backend.tensorflow_backend import _to_tensor, cast, flatten%0Aimport tensorflow as tf%0Adef image_categorical_crossentropy(y_true, y_pred, from_logits=False):%0A    if not from_logits:%0A        epsilon = _to_tensor(_EPSILON, y_pred.dtype.base_dtype)%0A        y_pred = tf.clip_by_value(y_pred, epsilon, 1 - epsilon)%0A        y_pred = tf.log(y_pred)%0A%0A    output_shape = y_pred.get_shape()%0A    targets = cast(flatten(y_true), 'int64')%0A    logits = tf.reshape(y_pred, %5B-1, int(output_shape%5B-1%5D)%5D)%0A%0A    res = tf.nn.sparse_softmax_cross_entropy_with_logits(%0A            labels=targets,%0A            logits=logits)%0A%0A    res = tf.reshape(res, tf.shape(y_pred)%5B:-1%5D)%0A%0A    return tf.reduce_sum(res, -1)
 %0A

Improve pixelcnn drastically