Currently, sample returns 2d-array shape of $[n_{samples} * n_{obs}, 1]$. I propose to "divide" the dimensions of samples and observations and return 3d-array shape of [n_samples, n_obs, 1] or in general case if input tensor have shape $[n_1, n_2, \dots, n_{m-1}, n_m]$ than output tensor will have shape $[n_1, n_2, \dots, n_{m-1}, n_{obs} ]$ or $[n_1, n_2, \dots, n_{m-1}, n_{obs}, 1]$ (see #6 and #5)
Related issues
#5 #6