How to use einops - 10 common examples

def test4(x):
        # space-to-depth
        y = rearrange(x, 'b c (h h1) (w w1) -> b (h1 w1 c) h w', h1=2, w1=2)
        assert y.shape == (10, 20 * 4, 30 // 2, 40 // 2)
        return y

def test8(x):
        # max-pooling
        y = reduce(x, 'b c (h h1) (w w1) -> b c h w', reduction='max', h1=2, w1=2)
        assert y.shape == (10, 20, 30 // 2, 40 // 2)
        return y

device = torch.device("cpu")

    model = VideoModel(pool_spatial=args.pool_spatial,
                       pool_temporal=args.pool_temporal)

    model.eval()

    for params in model.parameters():
        params.requires_grad = False

    model = model.to(device)
    model = nn.DataParallel(model)

    transform = Compose([
        ToTensor(),
        Rearrange("t h w c -> c t h w"),
        Resize(args.frame_size),
        Normalize(mean=[0.43216, 0.394666, 0.37645], std=[0.22803, 0.22145, 0.216989]),
    ])

    # dataset = WebcamDataset(clip=32, transform=transform)

    dataset = VideoDataset(args.video, clip=32, transform=transform)
    loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=0, shuffle=False)

    features = []

    with torch.no_grad():
        for inputs in tqdm(loader, total=len(dataset) // args.batch_size):
            inputs = inputs.to(device)

            outputs = model(inputs)

Observation = stft(observation)
    num_samples = observation.shape[-1]

    Y_mm = rearrange(Observation, 'd t f -> f t d')

    t = Trainer()
    affiliation = t.fit(
        Y_mm,
        num_classes=3,
        iterations=iterations * 2,
        weight_constant_axis=-1,
    ).predict(Y_mm)
    
    pa = DHTVPermutationAlignment.from_stft_size(512)
    affiliation_pa = pa(rearrange(affiliation, 'f k t -> k f t'))
    affiliation_pa = rearrange(affiliation_pa, 'k f t -> k t f')

    Speech_image_0_est, Speech_image_1_est, Noise_image_est = Observation[reference_channel, :, :] * affiliation_pa

    speech_image_0_est = istft(Speech_image_0_est, num_samples=num_samples)
    speech_image_1_est = istft(Speech_image_1_est, num_samples=num_samples)
    noise_image_est = istft(Noise_image_est, num_samples=num_samples)

    ###########################################################################
    # Calculate the metrics

    speech_image = ex['audio_data']['speech_image']
    noise_image = ex['audio_data']['noise_image']
    speech_source = ex['audio_data']['speech_source']

    Speech_image = stft(speech_image)
    Noise_image = stft(noise_image)

def new_way(input, num_classes, num_anchors, anchors, stride_h, stride_w):
        raw_predictions = rearrange(input, ' b (anchor prediction) h w -> prediction b anchor h w', anchor=num_anchors)

        anchors = torch.FloatTensor(anchors).to(input.device)
        anchor_sizes = rearrange(anchors, 'anchor dim -> dim () anchor () ()')

        _, _, _, in_h, in_w = raw_predictions.shape
        grid_h = rearrange(torch.arange(in_h).float(), 'h -> () () h ()').to(input.device)
        grid_w = rearrange(torch.arange(in_w).float(), 'w -> () () () w').to(input.device)

        predicted_bboxes = torch.zeros_like(raw_predictions)
        predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h  # center y
        predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w  # center x
        predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes  # bbox width and height
        predicted_bboxes[4] = raw_predictions[4].sigmoid()  # confidence
        predicted_bboxes[5:] = raw_predictions[5:].sigmoid()  # class predictions
        # only to match results of original code, not needed
        return rearrange(predicted_bboxes, 'prediction b anchor h w -> b anchor h w prediction')

for n_arrays in [1, 2, 5]:
            shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
            for shape in shapes:
                if backend.framework_name == 'mxnet.ndarray' and len(shape) == 0:
                    # known bug of mxnet
                    continue
                arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)]
                arrays2 = [backend.from_numpy(array) for array in arrays1]
                result0 = numpy.asarray(arrays1)
                result1 = rearrange(arrays1, '...->...')
                result2 = rearrange(arrays2, '...->...')
                assert numpy.array_equal(result0, result1)
                assert numpy.array_equal(result1, backend.to_numpy(result2))

                result1 = rearrange(arrays1, 'b ... -> ... b')
                result2 = rearrange(arrays2, 'b ... -> ... b')
                assert numpy.array_equal(result1, backend.to_numpy(result2))

anchors = torch.FloatTensor(anchors).to(input.device)
        anchor_sizes = rearrange(anchors, 'anchor dim -> dim () anchor () ()')

        _, _, _, in_h, in_w = raw_predictions.shape
        grid_h = rearrange(torch.arange(in_h).float(), 'h -> () () h ()').to(input.device)
        grid_w = rearrange(torch.arange(in_w).float(), 'w -> () () () w').to(input.device)

        predicted_bboxes = torch.zeros_like(raw_predictions)
        predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h  # center y
        predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w  # center x
        predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes  # bbox width and height
        predicted_bboxes[4] = raw_predictions[4].sigmoid()  # confidence
        predicted_bboxes[5:] = raw_predictions[5:].sigmoid()  # class predictions
        # only to match results of original code, not needed
        return rearrange(predicted_bboxes, 'prediction b anchor h w -> b anchor h w prediction')

def unet_like_1d(x, usual_convolution):
        # u-net like steps for increasing / reducing dimensionality
        x = rearrange(x, 'b c t1 t2 -> b c (t1 t2)')  # reduce dimensionality
        y = rearrange(x, 'b c (t dt) -> b (dt c) t', dt=2)
        y = usual_convolution(y)
        x = x + rearrange(y, 'b (dt c) t -> b c (t dt)', dt=2)
        return x

def test9(x):
        # squeeze - unsqueeze
        y = reduce(x, 'b c h w -> b c () ()', reduction='max')
        assert y.shape == (10, 20, 1, 1)
        y = rearrange(y, 'b c () () -> c b')
        assert y.shape == (20, 10)
        return y

def test_concatenations_and_stacking():
    for backend in imp_op_backends:
        print('testing shapes for ', backend.framework_name)
        for n_arrays in [1, 2, 5]:
            shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
            for shape in shapes:
                if backend.framework_name == 'mxnet.ndarray' and len(shape) == 0:
                    # known bug of mxnet
                    continue
                arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)]
                arrays2 = [backend.from_numpy(array) for array in arrays1]
                result0 = numpy.asarray(arrays1)
                result1 = rearrange(arrays1, '...->...')
                result2 = rearrange(arrays2, '...->...')
                assert numpy.array_equal(result0, result1)
                assert numpy.array_equal(result1, backend.to_numpy(result2))

                result1 = rearrange(arrays1, 'b ... -> ... b')
                result2 = rearrange(arrays2, 'b ... -> ... b')
                assert numpy.array_equal(result1, backend.to_numpy(result2))