@@ -0,0 +1,554 @@
+from __future__ import unicode_literals%0Aimport pytest%0Aimport pickle%0A%0Afrom ...api import with_flatten%0Afrom ...v2v import Affine%0A%0A@pytest.fixture%0Adef affine():%0A    return Affine(5, 3)%0A%0Adef test_pickle_with_flatten(affine):%0A    Xs = %5Baffine.ops.allocate((2, 3)),%0A          affine.ops.allocate((4, 3))%5D%0A    model = with_flatten(affine)%0A    pickled = pickle.dumps(model)%0A    loaded = pickle.loads(pickled)%0A    Ys = loaded(Xs)%0A    assert len(Ys) == 2%0A    assert Ys%5B0%5D.shape == (Xs%5B0%5D.shape%5B0%5D, affine.nO)%0A    assert Ys%5B1%5D.shape == (Xs%5B1%5D.shape%5B0%5D, affine.nO)%0A%0A

Add test for pickle