Method flatten_ref

caffe2/python/operator_test/flatten_op_test.py:24–28

(X)

                axis=axis)

            def flatten_ref(X):
                shape = X.shape
                outer = np.prod(shape[:axis]).astype(int)
                inner = np.prod(shape[axis:]).astype(int)
                return np.copy(X).reshape(outer, inner),

            self.assertReferenceChecks(gc, op, [X], flatten_ref)