def test_identity():
model = nn.Identity()
x = mx.np.random.uniform(size=(128, 33, 64))
assert_almost_equal(model(x), x)
def __init__(self, input_num, dim, **kwargs):
super(SingleConcat, self).__init__(**kwargs)
self.concat = nn.HybridConcatenate(axis=dim)
for i in range(input_num):
self.concat.add(nn.Identity())