def test_take(self):
a0 = numpy.arange(24.).reshape(4, 3, 2)
a1 = ctf.astensor(a0)
self.assertEqual(
ctf.take(a0, numpy.array([0, 3]), axis=0).shape, (2, 3, 2))
self.assertEqual(ctf.take(a1, [2], axis=1).shape, (4, 1, 2))
self.assertEqual(a1.take([0], axis=-1).shape, (4, 3, 1))
def test_take(self):
a0 = numpy.arange(24.).reshape(4,3,2)
a1 = ctf.astensor(a0)
self.assertEqual(ctf.take(a0, numpy.array([0,3]), axis=0).shape, (2,3,2))
self.assertEqual(ctf.take(a1, [2], axis=1).shape, (4,1,2))
self.assertEqual(a1.take([0], axis=-1).shape, (4,3,1))
self.assertEqual(a1.transpose().take([1], axis=0).shape, (1,3,4))