Files
2026-07-13 13:30:25 +08:00

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()