YannDubs

%matplotlib widget # need to install ipympl

import numpy as np
from matplotlib.lines import Line2D
from matplotlib.artist import Artist
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

import numpy as np

YannDubs

def stratify_sampling(x, n_samples, stratify):
    """Perform stratify sampling of a tensor.
    
    parameters
    ----------
    x: np.ndarray or torch.Tensor
        Array to sample from. Sampels from first dimension.
        
    n_samples: int
        Number of samples to sample