//  [python]
def get_generalization_in_dict(self, spatial_img, img, idx):
        # case where we use the convolutional encoder for generalization, either testing or training
        if (
            self.generalization_mode == "conv_cnp"
            or self.generalization_mode == "conv_cnp_test"
        ):
            if self.test_sparsity == "full":
                img_sparse = spatial_img
            elif self.test_sparsity == "half":
                img_sparse = spatial_img
                img_sparse[:, 16:, :] = 0.0
            else:
                if self.generalization_mode == "conv_cnp_test":
                    num_context = int(self.test_sparsity)
                else:
                    num_context = int(
                        torch.empty(1)
                        .uniform_(
                            self.train_sparsity_range[0], self.train_sparsity_range[1]
                        )
                        .item()
                    )
                mask = spatial_img.new_empty(
                    1, spatial_img.size(1), spatial_img.size(2)
                ).bernoulli_(p=num_context / np.prod(self.sidelength))
                img_sparse = mask * spatial_img
            in_dict = {"idx": idx, "coords": self.mgrid, "img_sparse": img_sparse}
        # case where we use the set encoder for generalization, either testing or training
        elif (
            self.generalization_mode == "cnp" or self.generalization_mode == "cnp_test"
        ):
            if self.test_sparsity == "full":
                in_dict = {
                    "idx": idx,
                    "coords": self.mgrid,
                    "img_sub": img,
                    "coords_sub": self.mgrid,
                }
            elif self.test_sparsity == "half":
                in_dict = {
                    "idx": idx,
                    "coords": self.mgrid,
                    "img_sub": img[:512, :],
                    "coords_sub": self.mgrid[:512, :],
                }
            else:
                if self.generalization_mode == "cnp_test":
                    subsamples = int(self.test_sparsity)
                    rand_idcs = np.random.choice(
                        img.shape[0], size=subsamples, replace=False
                    )
                    img_sparse = img[rand_idcs, :]
                    coords_sub = self.mgrid[rand_idcs, :]
                    in_dict = {
                        "idx": idx,
                        "coords": self.mgrid,
                        "img_sub": img_sparse,
                        "coords_sub": coords_sub,
                    }
                else:
                    subsamples = np.random.randint(
                        self.train_sparsity_range[0], self.train_sparsity_range[1]
                    )
                    rand_idcs = np.random.choice(
                        img.shape[0], size=self.train_sparsity_range[1], replace=False
                    )
                    img_sparse = img[rand_idcs, :]
                    coords_sub = self.mgrid[rand_idcs, :]

                    rand_idcs_2 = np.random.choice(
                        img_sparse.shape[0], size=subsamples, replace=False
                    )
                    ctxt_mask = torch.zeros(img_sparse.shape[0], 1)
                    ctxt_mask[rand_idcs_2, 0] = 1.0

                    in_dict = {
                        "idx": idx,
                        "coords": self.mgrid,
                        "img_sub": img_sparse,
                        "coords_sub": coords_sub,
                        "ctxt_mask": ctxt_mask,
                    }
        else:
            in_dict = {"idx": idx, "coords": self.mgrid}

        return in_dict