def __init__(self, nptype):
super(PackNet, self).__init__()
self.pack = P.Pack(axis=2)
self.data_np = np.array([0] * 16).astype(nptype)
self.data_np = np.reshape(self.data_np, (2, 2, 2, 2))
self.x1 = Parameter(initializer(
Tensor(self.data_np), [2, 2, 2, 2]), name='x1')
self.x2 = Parameter(initializer(
Tensor(np.arange(16).reshape(2, 2, 2, 2).astype(nptype)), [2, 2, 2, 2]), name='x2')
def pack_pynative(nptype):
context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
x1 = np.array([0] * 16).astype(nptype)
x1 = np.reshape(x1, (2, 2, 2, 2))
x1 = Tensor(x1)
x2 = Tensor(np.arange(16).reshape(2, 2, 2, 2).astype(nptype))
expect = np.array([[[[[0, 0], [0, 1]], [[0, 0], [2, 3]]],
[[[0, 0], [4, 5]], [[0, 0], [6, 7]]]],
[[[[0, 0], [8, 9]], [[0, 0], [10, 11]]],
[[[0, 0], [12, 13]], [[0, 0], [14,
15]]]]]).astype(nptype)
output = P.Pack(axis=2)((x1, x2))
assert (output.asnumpy() == expect).all()