Repository: lisa-lab/pylearn2

Issue:
Bug in FlattenerLayer.get_layer_monitoring_channels
Problem is that when raw_space in this function is composite space:
raw_space CompositeSpace(Conv2DSpace(shape=(4, 4), num_channels=64, axes=('c', 0, 1, 'b'), dtype=float32), Conv2DSpace(shape=(4, 4), num_channels=64, axes=('c', 0, 1, 'b'), dtype=float32), Conv2DSpace(shape=(4, 4), num_channels=64, axes=('c', 0, 1, 'b'), dtype=float32), Conv2DSpace(shape=(4, 4), num_channels=64, axes=('c', 0, 1, 'b'), dtype=float32), Conv2DSpace(shape=(4, 4), num_channels=64, axes=('c', 0, 1, 'b'), dtype=float32))
and raw_state:
raw_state (Reshape{2}.0, Reshape{2}.0, Reshape{2}.0, Reshape{2}.0, Reshape{2}.0)
and inside function when you execute code line:
raw_space.validate(raw_state)
you get error:
ValueError: The value of a Conv2DSpace batch must be 4D, got 2 dimensions for Reshape{2}.0.
Bug started with commit: cb6b2510af7f607f729822f8bddc54cace5864cf



Discussion:
It looks like I created another bug with that PR. Although, I'm pretty sure the "fix" I took from FlattenerLayer.cost, was originally a bug in your case (CompositeLayer with convolutional layers), if you let the FlattenerLayer/CompositeLayer get the default cost from FlattenerLayer.
https://github.com/lisa-lab/pylearn2/blob/master/pylearn2/models/mlp.py#L4387

The current code in both get_layer_monitoring_channels and cost doesn't try and reshape the parts pulled out of the Join which is a problem if the raw_layer contained convolutional layers.

I have a version working for all cases I've tested so far that I'll send in a PR today. The important part is below in the for loop. @lamblin was looking through the first PR.

```
if isinstance(raw_space, CompositeSpace):
             # Pick apart the Join that fprop used to make state.
             assert hasattr(state, 'owner')
             owner = state.owner
             assert owner is not None
             assert str(owner.op) == 'Join'
             # First input to join op is the axis.
             raw_state = tuple(owner.inputs[1:])
             state = ()
             for st,sp in safe_izip(raw_state, raw_space.components):
                 dim = sp.get_total_dimension()
                 state += (VectorSpace(dim).format_as(st, sp),)
             raw_space.validate(state)
         else:
             # Format state as layer output space.
             state = self.get_output_space().format_as(state, raw_space)
```

The FlattenerLayer.cost (which is where I borrowed the Join splitting code) breaks in a network like this with the same error as described in this issue. I think the code I posted above will also fix this, but I haven't tried it yet.

```
mlp = MLP(
         layers=[
             FlattenerLayer(
             CompositeLayer(
                     'composite',
                     [ConvElemwise(16,[2,2], 'sf1', SigmoidConvNonlinearity(), irange=0.1),
                     ConvElemwise(16,[2,2], 'sf2', SigmoidConvNonlinearity(), irange=0.1)]
             )
             )
         ],
         input_space=Conv2DSpace(shape=[10,10],num_channels=2)
         )

dataset = DenseDesignMatrix(
             topo_view=np.random.rand(20, 10,10,2).astype(theano.config.floatX),
             y=np.random.rand(20, 10).astype(theano.config.floatX))

train = Train(dataset, mlp, SGD(0.1, batch_size=5, monitoring_dataset=dataset))
train.algorithm.termination_criterion = EpochCounter(1)
train.main_loop()
```

The model above fails at:
https://github.com/lisa-lab/pylearn2/blob/master/pylearn2/expr/nnet.py#L155
if the space formatting is done because the op is a dimshuffle.

@JesseLivezey Your changes make sense, but regarding the crash in `kl()`, we would probably need to do one of the following:
- implement an `undo_format` method in `Space`s, or change the implementation of `format_as` to remove variables from the transformation graph, rather than adding an inverse transformation,
- update `kl` (or `SigmoidConvNonLinearity.cost`, maybe) to dig through `dimshuffle` operations, and apply the inverse transformation on the target instead,
- update `kl` to support an input that does not come from `sigmoid`, and check that Theano will automatically compile the resulting expression in the expression optimized for stability,
- for `get_layer_monitoring_channels`, we could use `state_below`, rather than try to undo what `fprop` did, but that would not solve the `cost` problem.


Resolving patch:
diff --git a/pylearn2/models/mlp.py b/pylearn2/models/mlp.py
--- a/pylearn2/models/mlp.py
+++ b/pylearn2/models/mlp.py
@@ -4335,20 +4335,8 @@ def get_layer_monitoring_channels(self, state_below=None,
                                       state=None, targets=None):
 
         raw_space = self.raw_layer.get_output_space()
-
-        if isinstance(raw_space, CompositeSpace):
-            # Pick apart the Join that fprop used to make state.
-            assert hasattr(state, 'owner')
-            owner = state.owner
-            assert owner is not None
-            assert str(owner.op) == 'Join'
-            # First input to join op in the axis.
-            raw_state = tuple(owner.inputs[1:])
-            raw_space.validate(raw_state)
-            state = raw_state
-        else:
-            # Format state as layer output space.
-            state = self.get_output_space().format_as(state, raw_space)
+        state = raw_space.undo_format_as(state,
+                                         self.get_output_space())
 
         if targets is not None:
             targets = self.get_target_space().format_as(
@@ -4404,22 +4392,9 @@ def cost(self, Y, Y_hat):
 
         raw_space = self.raw_layer.get_output_space()
         target_space = self.output_space
-        raw_Y = target_space.format_as(Y, raw_space)
 
-        if isinstance(raw_space, CompositeSpace):
-            # Pick apart the Join that our fprop used to make Y_hat
-            assert hasattr(Y_hat, 'owner')
-            owner = Y_hat.owner
-            assert owner is not None
-            assert str(owner.op) == 'Join'
-            # first input to join op is the axis
-            raw_Y_hat = tuple(owner.inputs[1:])
-        else:
-            # To implement this generally, we'll need to give Spaces an
-            # undo_format or something. You can't do it with format_as
-            # in the opposite direction because Layer.cost needs to be
-            # able to assume that Y_hat is the output of fprop
-            raise NotImplementedError()
+        raw_Y = target_space.format_as(Y, raw_space)
+        raw_Y_hat = raw_space.undo_format_as(Y_hat, target_space)
         raw_space.validate(raw_Y_hat)
 
         return self.raw_layer.cost(raw_Y, raw_Y_hat)
diff --git a/pylearn2/space/__init__.py b/pylearn2/space/__init__.py
--- a/pylearn2/space/__init__.py
+++ b/pylearn2/space/__init__.py
@@ -224,6 +224,34 @@ def _cast(arg, dtype):
         raise TypeError("Unsupported arg type '%s'" % str(type(arg)))
 
 
+def _undo_op(arg, string, strict=False):
+    """
+    Undo symbolic op if string is in str(op).
+
+    Returns <arg> untouched if there was no symbolic op.
+
+    Parameters
+    ----------
+    arg : any symbolic variable.
+    string : str
+        String that specifies op.
+    strict : bool
+        Whether to force op undo or not (default False).
+    """
+
+    if hasattr(arg.owner, 'op'):
+        owner = arg.owner
+        if string in str(owner.op):
+            return owner.inputs[0]
+        elif strict:
+            raise ValueError(string + ' not found in op ' +
+                             str(owner.op) + '.')
+    elif strict:
+        raise ValueError(string + ' op not found in variable ' +
+                         str(arg) + '.')
+    return arg
+
+
 class Space(object):
     """
     A vector space that can be transformed by a linear operator.
@@ -548,6 +576,88 @@ def _format_as_impl(self, is_numeric, batch, target_space):
         raise NotImplementedError("%s does not implement _format_as_impl()." %
                                   type(self))
 
+    def undo_np_format_as(self, batch, space):
+        """
+        Returns a numeric batch (e.g. a numpy.ndarray or scipy.sparse sparse
+        array), with formatting from space undone.
+
+        This is just a wrapper around self._undo_format_as, with an extra check
+        to throw an exception if <batch> is symbolic.
+
+        Parameters
+        ----------
+        batch : numpy.ndarray, or one of the scipy.sparse matrices.
+            Array which lies in this space.
+        space : Space
+            Space to undo formatting from.
+
+        Returns
+        -------
+        numpy.ndarray or one of the scipy.sparse matrices
+            The formatted batch.
+        """
+
+        self._check_is_numeric(batch)
+
+        return space.np_format_as(batch=batch,
+                                  space=self)
+
+    def undo_format_as(self, batch, space):
+        """
+        Returns a symbolic batch (e.g. a theano.tensor or theano.sparse
+        array), with formatting from space undone.
+
+        This is just a wrapper around self._undo_format_as, with an extra check
+        to throw an exception if <batch> is symbolic. Formatting to space
+
+        Parameters
+        ----------
+        batch : numpy.ndarray, or one of the scipy.sparse matrices.
+            Array which lies in this space.
+        space : Space
+            Space to undo formatting from.
+
+        Returns
+        -------
+        A symbolic Theano variable
+            The batch formatted as self.
+        """
+        self._check_is_symbolic(batch)
+        space.validate(batch)
+        self._check_sizes(space)
+
+        batch = self._undo_format_as_impl(batch=batch,
+                                          space=space)
+        # Checks if batch belongs to this space
+        self.validate(batch)
+
+        return batch
+
+    def _undo_format_as_impl(self, batch, target_space):
+        """
+        Actual implementation of undo_format_as.
+        Undoes target_space_formatting.
+        Note that undo_np_format_as calls np_format_as.
+
+        Parameters
+        ----------
+        batch : a theano symbol, or a nested tuple thereof
+            Implementations of this method may assume that batch lies in
+            space (i.e. that it passed self._validate(batch) without throwing
+            an exception).
+        target_space : A Space subclass
+            The space to undo batch formatting from.
+
+        Returns
+        -------
+        A symbolic Theano variable
+            The batch, converted from target_space, back to self.
+        """
+
+        raise NotImplementedError("%s does not implement "
+                                  "_undo_format_as_impl()." %
+                                  type(self))
+
     def validate(self, batch):
         """
         Runs all validate_callbacks, then checks that batch lies in this space.
@@ -1200,6 +1310,110 @@ def is_sparse(batch):
 
         return _cast(result, dtype=to_type)
 
+    @functools.wraps(Space._undo_format_as_impl)
+    def _undo_format_as_impl(self, batch, space):
+
+        def is_sparse(batch):
+            return isinstance(batch, theano.sparse.SparseVariable)
+
+        if not isinstance(space, IndexSpace):
+            my_dimension = self.get_total_dimension()
+            other_dimension = space.get_total_dimension()
+            if my_dimension != other_dimension:
+                raise ValueError(str(self) + " with total dimension " +
+                                 str(my_dimension) +
+                                 " can't undo format a batch from " +
+                                 str(space) +
+                                 "because its total dimension is " +
+                                 str(other_dimension))
+
+        if isinstance(space, CompositeSpace):
+            if isinstance(batch, theano.sparse.SparseVariable):
+                warnings.warn('Undo formatting from a sparse VectorSpace to a '
+                              'CompositeSpace is currently (2 Jan 2014) a '
+                              'non-differentiable action. This is because it '
+                              'calls slicing operations on a sparse batch '
+                              '(e.g. "my_matrix[r:R, c:C]", which Theano does '
+                              'not yet have a gradient operator for. If '
+                              'autodifferentiation is reporting an error, '
+                              'this may be why. Formatting batch type %s '
+                              'from space %s to space %s' %
+                              (type(batch), self, space))
+
+            # Recursively try and find a non-Composite, non-Null space
+            # to extract underlying theano variable
+            def extract_vector_variable(composite_space, batch_tuple):
+                found = False
+                for sp, el in safe_zip(composite_space.components,
+                                       batch_tuple):
+                    dim = sp.get_total_dimension()
+                    if not isinstance(sp, NullSpace) and dim > 0:
+                        if isinstance(sp, CompositeSpace):
+                            var, found = extract_vector_variable(sp, el)
+                            var = var.owner.inputs[0]
+                        else:
+                            dummy_sp = VectorSpace(dim=dim,
+                                                   sparse=sp.sparse,
+                                                   dtype=sp.dtype
+                                                   )
+                            var = dummy_sp.undo_format_as(el, sp)
+                            found = True
+                    if found:
+                        break
+                return var, found
+
+            var, found = extract_vector_variable(space, batch)
+            batch = var
+
+            if not found:
+                raise TypeError("Could not find a valid space "
+                                "to undo format from in the "
+                                "CompositeSpace.")
+            else:
+                # Undo subtensor slice
+                owner = batch.owner
+                assert 'Subtensor' in str(owner.op)
+                batch = owner.inputs[0]
+
+        elif isinstance(space, Conv2DSpace):
+            if is_sparse(batch):
+                raise TypeError("Undo formatting a SparseVariable to a "
+                                "Conv2DSpace is not supported, since "
+                                "neither scipy nor Theano has sparse "
+                                "tensors with more than 2 dimensions. "
+                                "We need 4 dimensions to represent a "
+                                "Conv2DSpace batch")
+            # Check for cast
+            batch = _undo_op(batch, 'Cast')
+            # Undo axes shuffle
+            if space.axes != space.default_axes:
+                batch = _undo_op(batch, 'DimShuffle', strict=True)
+            # Undo reshape
+            batch = _undo_op(batch, 'Reshape{4}', strict=True)
+
+        elif isinstance(space, VectorSpace):
+            if self.dim != space.dim:
+                raise ValueError("Can't convert between VectorSpaces of "
+                                 "different sizes (%d to %d)."
+                                 % (self.dim, space.dim))
+            # Check for cast
+            batch = _undo_op(batch, 'Cast')
+            # Undo any sparse-dense switches
+            if self.sparse != is_sparse(batch):
+                if space.sparse:
+                    batch = _undo_op(batch, 'SparseFromDense', strict=True)
+                elif isinstance(batch, theano.sparse.SparseVariable):
+                    batch = _undo_op(batch, 'DenseFromSparse', strict=True)
+                else:
+                    assert False, ("Unplanned-for branch in if-elif "
+                                   "chain. This is a bug in the code.")
+
+        else:
+            raise NotImplementedError("%s doesn't know how to format as %s" %
+                                      (self, space))
+
+        return batch
+
     def __eq__(self, other):
         """
         .. todo::
@@ -1837,6 +2051,30 @@ def _format_as_impl(self, is_numeric, batch, space):
 
         return _cast(result, space.dtype)
 
+    @functools.wraps(Space._undo_format_as_impl)
+    def _undo_format_as_impl(self, batch, space):
+        # Check for cast
+        batch = _undo_op(batch, 'Cast')
+
+        if isinstance(space, VectorSpace):
+            # Check for SparseFromDense
+            batch = _undo_op(batch, 'SparseFromDense')
+            # Undo reshape op
+            batch = _undo_op(batch, 'Reshape', strict=True)
+            # Check to see if axis ordering was changed
+            if self.axes != self.default_axes:
+                batch = _undo_op(batch, 'DimShuffle', strict=True)
+
+        elif isinstance(space, Conv2DSpace):
+            # Check to see if axis ordering was changed
+            if space.axes != self.axes:
+                batch = _undo_op(batch, 'DimShuffle', strict=True)
+        else:
+            raise NotImplementedError("%s doesn't know how to format as %s"
+                                      % (str(self), str(space)))
+
+        return batch
+
 
 class CompositeSpace(Space):
     """
@@ -2083,6 +2321,105 @@ def recursive_format_as(orig_space, batch, dest_space):
                                   " does not know how to format as " +
                                   str(space))
 
+    @functools.wraps(Space._undo_format_as_impl)
+    def _undo_format_as_impl(self, batch, space):
+        """
+        Undoes the formatting to a single VectorSpace, or to a CompositeSpace.
+
+        CompositeSpace->VectorSpace:
+          Traverses the nested components in depth-first order, serializing the
+          leaf nodes (i.e. the non-composite subspaces) into the VectorSpace.
+
+        CompositeSpace->CompositeSpace:
+
+          Only works for two CompositeSpaces that have the same nested
+          structure. Traverses both CompositeSpaces' nested components in
+          parallel, converting between corresponding non-composite components
+          in <self> and <space> as:
+
+              `self_component._format_as(is_numeric,
+                                         batch_component,
+                                         space_component)`
+
+        Parameters
+        ----------
+        batch : WRITEME
+        space : WRITEME
+
+        Returns
+        -------
+        WRITEME
+        """
+
+        if isinstance(space, VectorSpace):
+            # Undo join
+            if space.sparse:
+                owner = batch.owner
+                assert owner is not None
+                assert 'HStack' in str(owner.op)
+                batch = owner.inputs
+            else:
+                owner = batch.owner
+                assert owner is not None
+                assert str(owner.op) == 'Join'
+                # First component is join axis
+                batch = owner.inputs[1:]
+
+            def extract_dtype(dtype):
+                if isinstance(dtype, tuple):
+                    return extract_dtype(dtype[0])
+                else:
+                    return dtype
+
+            def compose_batch(composite_space, batch_list):
+                rval = ()
+                for sp, bt in safe_zip(composite_space.components, batch_list):
+                    if False and isinstance(sp, CompositeSpace):
+                        composed, batch_list = compose_batch(sp, batch_list)
+                        rval += (composed,)
+                    else:
+                        sparse = getattr(sp, 'sparse', False)
+                        dtype = extract_dtype(sp.dtype)
+                        new_sp = VectorSpace(dim=sp.get_total_dimension(),
+                                             dtype=dtype,
+                                             sparse=sparse
+                                             )
+                        new_batch = sp.undo_format_as(bt,
+                                                      new_sp)
+                        rval += (new_batch,)
+
+                return rval
+            composed = compose_batch(self, batch)
+            return composed
+
+        if isinstance(space, CompositeSpace):
+            def recursive_undo_format_as(orig_space, batch, dest_space):
+                if not (isinstance(orig_space, CompositeSpace) ==
+                        isinstance(dest_space, CompositeSpace)):
+                    raise TypeError("Can't convert between CompositeSpaces "
+                                    "with different tree structures")
+
+                # No need to check batch's tree structure.
+                # Space.undo_format_as() already did that
+                # by calling _validate(), before calling this
+                # method.
+
+                if isinstance(orig_space, CompositeSpace):
+                    return tuple(recursive_undo_format_as(os, bt, ds)
+                                 for os, bt, ds
+                                 in safe_zip(orig_space.components,
+                                             batch,
+                                             dest_space.components))
+                else:
+                    return orig_space.undo_format_as(batch,
+                                                     dest_space)
+
+            return recursive_undo_format_as(self, batch, space)
+
+        raise NotImplementedError(str(self) +
+                                  " does not know how to format as " +
+                                  str(space))
+
     @functools.wraps(Space._validate_impl)
     def _validate_impl(self, is_numeric, batch):
         if not isinstance(batch, tuple):


Test patch:
diff --git a/pylearn2/models/tests/test_mlp.py b/pylearn2/models/tests/test_mlp.py
--- a/pylearn2/models/tests/test_mlp.py
+++ b/pylearn2/models/tests/test_mlp.py
@@ -18,7 +18,8 @@
 from pylearn2.models.mlp import (FlattenerLayer, MLP, Linear, Softmax, Sigmoid,
                                  exhaustive_dropout_average,
                                  sampled_dropout_average, CompositeLayer,
-                                 mean_pool)
+                                 mean_pool, SigmoidConvNonlinearity,
+                                 ConvElemwise)
 from pylearn2.space import VectorSpace, CompositeSpace, Conv2DSpace
 from pylearn2.utils import is_iterable, sharedX
 from pylearn2.expr.nnet import pseudoinverse_softmax_numpy
@@ -551,25 +552,41 @@ def test_flattener_layer_state_separation_for_softmax():
     Creates a CompositeLayer wrapping two Softmax layers
     and ensures that state gets correctly picked apart.
     """
-    mlp = MLP(
-        layers=[
-            FlattenerLayer(
-                CompositeLayer(
-                    'composite',
-                    [Softmax(5, 'sf1', 0.1),
-                     Softmax(5, 'sf2', 0.1)]
-                )
-            )
-        ],
-        nvis=2
-        )
+    soft1 = Softmax(5, 'sf1', .1)
+    soft2 = Softmax(5, 'sf2', .1)
+    mlp = MLP(layers=[FlattenerLayer(CompositeLayer('comp',
+                                                    [soft1, soft2]))],
+              nvis=2)
+
+    X = np.random.rand(20, 2).astype(theano.config.floatX)
+    y = np.random.rand(20, 10).astype(theano.config.floatX)
+    dataset = DenseDesignMatrix(X=X, y=y)
+
+    train = Train(dataset, mlp, SGD(0.1,
+                                    batch_size=5,
+                                    monitoring_dataset=dataset))
+    train.algorithm.termination_criterion = EpochCounter(1)
+    train.main_loop()
 
-    dataset = DenseDesignMatrix(
-        X=np.random.rand(20, 2).astype(theano.config.floatX),
-        y=np.random.rand(20, 10).astype(theano.config.floatX))
 
-    train = Train(dataset, mlp,
-                  SGD(0.1, batch_size=5, monitoring_dataset=dataset))
+def test_flattener_layer_state_separation_for_conv():
+    """
+    Creates a CompositeLayer wrapping two Conv layers
+    and ensures that state gets correctly picked apart.
+    """
+    conv1 = ConvElemwise(8, [2, 2], 'sf1', SigmoidConvNonlinearity(), .1)
+    conv2 = ConvElemwise(8, [2, 2], 'sf2', SigmoidConvNonlinearity(), .1)
+    mlp = MLP(layers=[FlattenerLayer(CompositeLayer('comp',
+                                                    [conv1, conv2]))],
+              input_space=Conv2DSpace(shape=[5, 5], num_channels=2))
+
+    topo_view = np.random.rand(10, 5, 5, 2).astype(theano.config.floatX)
+    y = np.random.rand(10, 256).astype(theano.config.floatX)
+    dataset = DenseDesignMatrix(topo_view=topo_view, y=y)
+
+    train = Train(dataset, mlp, SGD(0.1,
+                                    batch_size=5,
+                                    monitoring_dataset=dataset))
     train.algorithm.termination_criterion = EpochCounter(1)
     train.main_loop()
 
diff --git a/pylearn2/space/tests/test_space.py b/pylearn2/space/tests/test_space.py
--- a/pylearn2/space/tests/test_space.py
+++ b/pylearn2/space/tests/test_space.py
@@ -1128,3 +1128,237 @@ def make_composite_space(dtype0, dtype1, use_conv2d):
         for to_space in all_spaces:
             for is_numeric in (True, False):
                 test_format(from_space, to_space, is_numeric)
+
+
+def test_symbolic_undo_format_as():
+    """
+    Test functionality of undo_format_as on symbolic batches.
+    After format_as and undo_format_as, the theano variable
+    should be the same object, not just an equivalent
+    variable.
+    """
+
+    # Compare identity of Composite batches
+    def assert_components(batch1, batch2):
+        for e1, e2 in zip(batch1, batch2):
+            if isinstance(e1, tuple) and isinstance(e2, tuple):
+                assert_components(e1, e2)
+            elif isinstance(e1, tuple) or isinstance(e2, tuple):
+                raise ValueError('Composite batches do not match.')
+            else:
+                assert e1 is e2
+
+    # VectorSpace and Conv2DSpace
+    VS = VectorSpace(dim=27)
+    VS_sparse = VectorSpace(dim=27, sparse=True)
+
+    # VectorSpace to Sparse VectorSpace
+    VS_batch = VS.make_theano_batch()
+    new_SVS_batch = VS.format_as(VS_batch, VS_sparse)
+    new_VS_batch = VS.undo_format_as(new_SVS_batch, VS_sparse)
+    assert new_VS_batch is VS_batch
+    assert new_SVS_batch is not VS_batch
+
+    # ConvSpace to ConvSpace
+    CS = Conv2DSpace(shape=[3, 3],
+                     num_channels=3,
+                     axes=('b', 0, 1, 'c'),
+                     dtype='float32')
+    CS_non_default = Conv2DSpace(shape=[3, 3],
+                                 num_channels=3,
+                                 axes=('c', 'b', 0, 1),
+                                 dtype='float64')
+    CS_batch = CS.make_theano_batch()
+    new_ndCS_batch = CS.format_as(CS_batch, CS_non_default)
+    new_CS_batch = CS.undo_format_as(new_ndCS_batch, CS_non_default)
+    assert new_CS_batch is CS_batch
+    assert new_ndCS_batch is not CS_batch
+    assert new_ndCS_batch.dtype == 'float64'
+    assert new_CS_batch.dtype == 'float32'
+
+    ndCS_batch = CS_non_default.make_theano_batch()
+    new_CS_batch = CS_non_default.format_as(ndCS_batch, CS)
+    new_ndCS_batch = CS_non_default.undo_format_as(new_CS_batch, CS)
+    assert new_ndCS_batch is ndCS_batch
+    assert new_CS_batch is not ndCS_batch
+    assert new_ndCS_batch.dtype == 'float64'
+    assert new_CS_batch.dtype == 'float32'
+
+    # Start in VectorSpace
+    VS_batch = VS.make_theano_batch()
+    new_CS_batch = VS.format_as(VS_batch, CS)
+    new_VS_batch = VS.undo_format_as(new_CS_batch, CS)
+    assert new_VS_batch is VS_batch
+
+    new_CS_batch = VS.format_as(VS_batch, CS_non_default)
+    new_VS_batch = VS.undo_format_as(new_CS_batch, CS_non_default)
+    assert new_VS_batch is VS_batch
+
+    # Start in Conv2D with default axes
+    CS_batch = CS.make_theano_batch()
+    new_VS_batch = CS.format_as(CS_batch, VS)
+    new_CS_batch = CS.undo_format_as(new_VS_batch, VS)
+    assert new_CS_batch is CS_batch
+    # Non-default axes
+    CS_batch = CS_non_default.make_theano_batch()
+    new_VS_batch = CS_non_default.format_as(CS_batch, VS)
+    new_CS_batch = CS_non_default.undo_format_as(new_VS_batch, VS)
+    assert new_CS_batch is CS_batch
+
+    # Composite Space to VectorSpace
+    VS = VectorSpace(dim=27)
+    CS = Conv2DSpace(shape=[2, 2], num_channels=3, axes=('b', 0, 1, 'c'))
+    CompS = CompositeSpace((CompositeSpace((VS, VS)), CS))
+    VS_large = VectorSpace(dim=(2*27+12))
+    CompS_batch = CompS.make_theano_batch()
+    new_VS_batch = CompS.format_as(CompS_batch, VS_large)
+    new_CompS_batch = CompS.undo_format_as(new_VS_batch, VS_large)
+    assert_components(CompS_batch, new_CompS_batch)
+
+    # VectorSpace to Composite Space
+    CompS = CompositeSpace((CompositeSpace((VS, VS)), CS))
+    VS_batch = VS_large.make_theano_batch()
+    new_CompS_batch = VS_large.format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_format_as(new_CompS_batch, CompS)
+    assert VS_batch is new_VS_batch
+    # Reorder CompositeSpace
+    CompS = CompositeSpace((VS, CompositeSpace((VS, CS))))
+    VS_batch = VS_large.make_theano_batch()
+    new_CompS_batch = VS_large.format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_format_as(new_CompS_batch, CompS)
+    assert VS_batch is new_VS_batch
+    # Reorder CompositeSpace
+    CompS = CompositeSpace((CompositeSpace((CompositeSpace((VS,)), CS)), VS))
+    VS_batch = VS_large.make_theano_batch()
+    new_CompS_batch = VS_large.format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_format_as(new_CompS_batch, CompS)
+    assert VS_batch is new_VS_batch
+
+    # CompositeSpace to CompositeSpace
+    VS = VectorSpace(dim=27)
+    CS = Conv2DSpace(shape=[3, 3], num_channels=3, axes=('b', 0, 1, 'c'))
+    CompS_VS = CompositeSpace((CompositeSpace((VS, VS)), VS))
+    CompS_CS = CompositeSpace((CompositeSpace((CS, CS)), CS))
+    CompS_VS_batch = CompS_VS.make_theano_batch()
+    new_CompS_CS_batch = CompS_VS.format_as(CompS_VS_batch, CompS_CS)
+    new_CompS_VS_batch = CompS_VS.undo_format_as(new_CompS_CS_batch, CompS_CS)
+    assert_components(CompS_VS_batch, new_CompS_VS_batch)
+
+
+def test_numeric_undo_format_as():
+    """
+    Test functionality of undo_np_format_as on numeric batches.
+    This calls np_format_as with spaces reversed.
+    """
+
+    # Compare identity of Composite batches
+    def assert_components(batch1, batch2):
+        for e1, e2 in zip(batch1, batch2):
+            if isinstance(e1, tuple) and isinstance(e2, tuple):
+                assert_components(e1, e2)
+            elif isinstance(e1, tuple) or isinstance(e2, tuple):
+                raise ValueError('Composite batches do not match.')
+            else:
+                assert np.allclose(e1, e2)
+
+    # VectorSpace and Conv2DSpace
+    VS = VectorSpace(dim=27)
+    VS_sparse = VectorSpace(dim=27, sparse=True)
+
+    # VectorSpace to Sparse VectorSpace
+    VS_batch = np.arange(270).reshape(10, 27)
+    new_SVS_batch = VS.np_format_as(VS_batch, VS_sparse)
+    new_VS_batch = VS.undo_np_format_as(new_SVS_batch, VS_sparse)
+    assert np.allclose(new_VS_batch, VS_batch)
+
+    # ConvSpace to ConvSpace
+    CS = Conv2DSpace(shape=[3, 3],
+                     num_channels=3,
+                     axes=('b', 0, 1, 'c'),
+                     dtype='float32')
+    CS_non_default = Conv2DSpace(shape=[3, 3],
+                                 num_channels=3,
+                                 axes=('c', 'b', 0, 1),
+                                 dtype='float64')
+    CS_batch = np.arange(270).reshape(10, 3, 3, 3).astype('float32')
+    new_ndCS_batch = CS.np_format_as(CS_batch, CS_non_default)
+    new_CS_batch = CS.undo_np_format_as(new_ndCS_batch, CS_non_default)
+    assert np.allclose(new_CS_batch, CS_batch)
+    assert new_ndCS_batch.shape != CS_batch.shape
+    assert new_ndCS_batch.dtype == 'float64'
+    assert new_CS_batch.dtype == 'float32'
+
+    ndCS_batch = np.arange(270).reshape(3, 10, 3, 3)
+    new_CS_batch = CS_non_default.np_format_as(ndCS_batch, CS)
+    new_ndCS_batch = CS_non_default.undo_np_format_as(new_CS_batch, CS)
+    assert np.allclose(new_ndCS_batch, ndCS_batch)
+    assert new_CS_batch.shape != ndCS_batch.shape
+    assert new_ndCS_batch.dtype == 'float64'
+    assert new_CS_batch.dtype == 'float32'
+
+    # Start in VectorSpace
+    VS_batch = np.arange(270).reshape(10, 27)
+    new_CS_batch = VS.np_format_as(VS_batch, CS)
+    new_VS_batch = VS.undo_np_format_as(new_CS_batch, CS)
+    assert np.allclose(new_VS_batch, VS_batch)
+    # Non-default axes
+    new_CS_batch = VS.np_format_as(VS_batch, CS_non_default)
+    new_VS_batch = VS.undo_np_format_as(new_CS_batch, CS_non_default)
+    assert np.allclose(new_VS_batch, VS_batch)
+
+    # Start in Conv2D with default axes
+    CS_batch = np.arange(270).reshape(10, 3, 3, 3)
+    new_VS_batch = CS.np_format_as(CS_batch, VS)
+    new_CS_batch = CS.undo_np_format_as(new_VS_batch, VS)
+    assert np.allclose(new_CS_batch, CS_batch)
+    # Non-default axes
+    CS_batch = np.arange(270).reshape(3, 10, 3, 3)
+    new_VS_batch = CS_non_default.np_format_as(CS_batch, VS)
+    new_CS_batch = CS_non_default.undo_np_format_as(new_VS_batch, VS)
+    assert np.allclose(new_CS_batch, CS_batch)
+
+    # Composite Space to VectorSpace
+    VS = VectorSpace(dim=27)
+    CS = Conv2DSpace(shape=[2, 2], num_channels=3, axes=('b', 0, 1, 'c'))
+    CompS = CompositeSpace((CompositeSpace((VS, VS)), CS))
+    VS_large = VectorSpace(dim=(2*27+12))
+    VS_batch = np.arange(270).reshape(10, 27)
+    VS_batch2 = 2*np.arange(270).reshape(10, 27)
+    CS_batch = 3*np.arange(120).reshape(10, 2, 2, 3)
+    CompS_batch = ((VS_batch, VS_batch2), CS_batch)
+    new_VS_batch = CompS.np_format_as(CompS_batch, VS_large)
+    new_CompS_batch = CompS.undo_np_format_as(new_VS_batch, VS_large)
+    assert_components(CompS_batch, new_CompS_batch)
+
+    # VectorSpace to Composite Space
+    CompS = CompositeSpace((CompositeSpace((VS, VS)), CS))
+    VS_batch = np.arange((2*27+12)*10).reshape(10, 2*27+12)
+    new_CompS_batch = VS_large.np_format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_np_format_as(new_CompS_batch, CompS)
+    assert np.allclose(VS_batch, new_VS_batch)
+    # Reorder CompositeSpace
+    CompS = CompositeSpace((VS, CompositeSpace((VS, CS))))
+    VS_batch = np.arange((2*27+12)*10).reshape(10, 2*27+12)
+    new_CompS_batch = VS_large.np_format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_np_format_as(new_CompS_batch, CompS)
+    assert np.allclose(VS_batch, new_VS_batch)
+    # Reorder CompositeSpace
+    CompS = CompositeSpace((CompositeSpace((CompositeSpace((VS,)), CS)), VS))
+    VS_batch = np.arange((2*27+12)*10).reshape(10, 2*27+12)
+    new_CompS_batch = VS_large.np_format_as(VS_batch, CompS)
+    new_VS_batch = VS_large.undo_np_format_as(new_CompS_batch, CompS)
+    assert np.allclose(VS_batch, new_VS_batch)
+
+    # CompositeSpace to CompositeSpace
+    VS = VectorSpace(dim=27)
+    CS = Conv2DSpace(shape=[3, 3], num_channels=3, axes=('b', 0, 1, 'c'))
+    VS_batch = np.arange(270).reshape(10, 27)
+    VS_batch2 = 2*np.arange(270).reshape(10, 27)
+    VS_batch3 = 3*np.arange(270).reshape(10, 27)
+    CompS_VS = CompositeSpace((CompositeSpace((VS, VS)), VS))
+    CompS_CS = CompositeSpace((CompositeSpace((CS, CS)), CS))
+    CompS_VS_batch = ((VS_batch, VS_batch2), VS_batch3)
+    new_CompS_CS_batch = CompS_VS.np_format_as(CompS_VS_batch, CompS_CS)
+    new_CompS_VS_batch = CompS_VS.undo_np_format_as(new_CompS_CS_batch,
+                                                    CompS_CS)
+    assert_components(CompS_VS_batch, new_CompS_VS_batch)