wehub-resource-sync
20 lines
471 B
Python
Executable File
20 lines
471 B
Python
Executable File
import unittest
|
|
import numpy as np
|
|
from prml import nn
|
|
from prml.nn.array.broadcast import broadcast_to
|
|
|
|
|
|
class TestBroadcastTo(unittest.TestCase):
|
|
|
|
def test_broadcast(self):
|
|
x = nn.Parameter(np.ones((1, 1)))
|
|
shape = (5, 2, 3)
|
|
y = broadcast_to(x, shape)
|
|
self.assertEqual(y.shape, shape)
|
|
y.backward(np.ones(shape))
|
|
self.assertTrue((x.grad == np.ones((1, 1)) * 30).all())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|