def test_transpose(): shapes = [(2, 3), (3, 1), (2, 3, 4)] for shape in shapes: x_np = np.random.random(shape) x_jp = jp.array(x_np) y_np = np.transpose(x_np) y_jp = jp.transpose(x_jp) y_jp = y_jp.numpy() assert y_jp.shape == y_np.shape
def test_permute(): shapes = [[(2, 3), [0, 1], [1, 0]], [(2, 1), [0, 1], [1, 0]], [(2, 3, 4), [0, 1, 2], [0, 2, 1], [1, 0, 2], [1, 2, 0], [2, 0, 1], [2, 1, 0]]] for shape in shapes: x_np = np.random.random(shape[0]) x_jp = jp.array(x_np) for dims in shape[1:]: y_np = np.transpose(x_np, dims) y_jp = jp.transpose(x_jp, dims) assert y_jp.shape == y_np.shape
def test_permute(): shapes = [] shapes.append([(2, 3), [0, 1], [1, 0]]) shapes.append([(2, 1), [0, 1], [1, 0]]) shapes.append([(2, 3, 4), [0, 1, 2], [0, 2, 1], [1, 0, 2], [1, 2, 0], [2, 0, 1], [2, 1, 0]]) for shape in shapes: x_np = np.random.random(shape[0]) x_jp = jp.array(x_np) for dims in shape[1:]: y_np = np.transpose(x_np, dims) y_jp = jp.transpose(x_jp, dims) assert y_jp.shape == y_np.shape