22 lines
524 B
Python
Executable File
22 lines
524 B
Python
Executable File
import unittest
|
|
import numpy as np
|
|
from prml import nn
|
|
|
|
|
|
class TestFlatten(unittest.TestCase):
|
|
|
|
def test_flatten(self):
|
|
self.assertRaises(TypeError, nn.flatten, "abc")
|
|
self.assertRaises(ValueError, nn.flatten, np.ones(1))
|
|
|
|
x = np.random.rand(5, 4)
|
|
p = nn.Parameter(x)
|
|
y = p.flatten()
|
|
self.assertTrue((y.value == x.flatten()).all())
|
|
y.backward(np.ones(20))
|
|
self.assertTrue((p.grad == np.ones((5, 4))).all())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|