New style flatten

chainer · Aug 24, 2017 · f6543af · f6543af
1 parent 7930acf
commit f6543af
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 17 deletions.
diff --git a/chainer/functions/__init__.py b/chainer/functions/__init__.py
@@ -47,7 +47,6 @@
 from chainer.functions.array.expand_dims import expand_dims  # NOQA
 from chainer.functions.array.expand_dims import ExpandDims  # NOQA
 from chainer.functions.array.flatten import flatten  # NOQA
-from chainer.functions.array.flatten import Flatten  # NOQA
 from chainer.functions.array.fliplr import fliplr  # NOQA
 from chainer.functions.array.fliplr import FlipLR  # NOQA
 from chainer.functions.array.flipud import flipud  # NOQA

diff --git a/chainer/functions/array/flatten.py b/chainer/functions/array/flatten.py
@@ -1,17 +1,4 @@
-from chainer import function
-
-
-class Flatten(function.Function):
-
-    """Flatten function."""
-
-    def forward(self, inputs):
-        self.retain_inputs(())
-        self._in_shape = inputs[0].shape
-        return inputs[0].ravel(),
-
-    def backward(self, inputs, grads):
-        return grads[0].reshape(self._in_shape),
+from chainer.functions.array import reshape
 
 
 def flatten(x):
@@ -50,4 +37,4 @@ def flatten(x):
         array([0, 1, 2, 3, 4, 5, 6, 7])
 
     """
-    return Flatten()(x)
+    return reshape.reshape(x, (x.size,))
diff --git a/tests/chainer_tests/functions_tests/array_tests/test_flatten.py b/tests/chainer_tests/functions_tests/array_tests/test_flatten.py
@@ -40,7 +40,7 @@ def test_forward_gpu(self):
 
     def check_backward(self, x_data, g_data):
         gradient_check.check_backward(
-            functions.Flatten(), x_data, g_data, dtype=numpy.float64)
+            functions.flatten, x_data, g_data, dtype=numpy.float64)
 
     def test_backward_cpu(self):
         self.check_backward(self.x, self.g)