Spaces:
Runtime error
Runtime error
File size: 1,069 Bytes
19c4ddf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from typing import Iterable, Union
import numpy as np
import torch
ArrayType = Union[np.ndarray, Iterable[int], torch.Tensor]
def to_torch(arr: ArrayType, dtype=torch.float):
if isinstance(arr, torch.Tensor):
return arr
return torch.from_numpy(np.array(arr)).to(dtype)
def sample_pmf(pmf: torch.Tensor, n_samples: int) -> torch.Tensor:
"""
Sample from the given discrete probability distribution with replacement.
The i-th bin is assumed to have mass pmf[i].
:param pmf: [batch_size, *shape, n_samples, 1] where (pmf.sum(dim=-2) == 1).all()
:param n_samples: number of samples
:return: indices sampled with replacement
"""
*shape, support_size, last_dim = pmf.shape
assert last_dim == 1
cdf = torch.cumsum(pmf.view(-1, support_size), dim=1)
inds = torch.searchsorted(cdf, torch.rand(cdf.shape[0], n_samples, device=cdf.device))
return inds.view(*shape, n_samples, 1).clamp(0, support_size - 1)
def safe_divide(a, b, epsilon=1e-6):
return a / torch.where(b < 0, b - epsilon, b + epsilon)
|