from __future__ import print_function
import plac

from thinc import datasets
from thinc.base import Network
from thinc.vec2vec import ReLu, Softmax
from thinc.util import score_model


class ReLuMLP(Network):
    Hidden = ReLu
    Output = Softmax
    width = 64
    depth = 3

    def setup(self, nr_out, nr_in, **kwargs):
        for i in range(self.depth):
            self.layers.append(self.Hidden(self.width, nr_in))
            nr_in = self.width
        self.layers.append(self.Output(nr_out, nr_in))


def main(batch_size=128, nb_epoch=8, nb_classes=10):
    model = ReLuMLP(10, 784)
    train_data, check_data, test_data = datasets.keras_mnist()
    
    with model.begin_training(train_data) as (trainer, optimizer):
        for examples, truth in trainer.iterate(model, train_data, check_data):
            guess, update = model.begin_update(examples)

            gradient, loss = model.get_gradient(guess, truth)
            optimizer.set_loss(loss)

            finish_update(gradient, optimizer, L2=trainer.L2)

    print('Test score:', score_model(model, test_data))


if __name__ == '__main__':
    if 1:
        plac.call(main)
    else:
        import cProfile
        import pstats
        cProfile.runctx("main()", globals(), locals(), "Profile.prof")
        s = pstats.Stats("Profile.prof")
        s.strip_dirs().sort_stats("time").print_stats()

@@ -44,16 +44,35 @@
 t plac%0A%0A
+import numpy as np%0A
 from thi
@@ -416,16 +416,23 @@
 .Hidden(
+nr_out=
 self.wid
@@ -435,16 +435,22 @@
 .width, 
+nr_in=
 nr_in))%0A
@@ -529,17 +529,203 @@
 _out
-, nr_in))
+=nr_out, nr_in=nr_in))%0A%0Adef get_gradient(scores, labels):%0A    target = np.zeros(scores.shape)%0A    for i, label in enumerate(labels):%0A        target%5Bi, int(label)%5D = 1.0%0A    return scores - target
 %0A%0A%0Ad
@@ -1035,41 +1035,46 @@
     
-guess, update = model.begin_updat
+assert hasattr(examples, 'shape'), typ
 e(ex
@@ -1081,17 +1081,16 @@
 amples)%0A
-%0A
         
@@ -1098,91 +1098,56 @@
    g
-radient, loss = model.get_gradient(guess, truth)%0A            optimizer.set_loss(los
+uess, finish_update = model.begin_update(example
 s)%0A%0A
@@ -1172,24 +1172,28 @@
 _update(
+get_
 gradient
 , optimi
@@ -1188,34 +1188,33 @@
 ient
-, optimizer, L2=trainer.L2
+(guess, truth), optimizer
 )%0A%0A 

Fix mnist example script