| |
| |
| |
|
|
| import numpy as np |
| import torch |
|
|
| __all__ = [ |
| "torch_randint", |
| "torch_random", |
| "torch_shuffle", |
| "torch_uniform", |
| "torch_random_choices", |
| ] |
|
|
|
|
| def torch_randint( |
| low: int, high: int, generator: torch.Generator or None = None |
| ) -> int: |
| """uniform: [low, high)""" |
| if low == high: |
| return low |
| else: |
| assert low < high |
| return int(torch.randint(low=low, high=high, generator=generator, size=(1,))) |
|
|
|
|
| def torch_random(generator: torch.Generator or None = None) -> float: |
| """uniform distribution on the interval [0, 1)""" |
| return float(torch.rand(1, generator=generator)) |
|
|
|
|
| def torch_shuffle( |
| src_list: list[any], generator: torch.Generator or None = None |
| ) -> list[any]: |
| rand_indexes = torch.randperm(len(src_list), generator=generator).tolist() |
| return [src_list[i] for i in rand_indexes] |
|
|
|
|
| def torch_uniform( |
| low: float, high: float, generator: torch.Generator or None = None |
| ) -> float: |
| """uniform distribution on the interval [low, high)""" |
| rand_val = torch_random(generator) |
| return (high - low) * rand_val + low |
|
|
|
|
| def torch_random_choices( |
| src_list: list[any], |
| generator: torch.Generator or None = None, |
| k=1, |
| weight_list: list[float] or None = None, |
| ) -> any or list: |
| if weight_list is None: |
| rand_idx = torch.randint( |
| low=0, high=len(src_list), generator=generator, size=(k,) |
| ) |
| out_list = [src_list[i] for i in rand_idx] |
| else: |
| assert len(weight_list) == len(src_list) |
| accumulate_weight_list = np.cumsum(weight_list) |
|
|
| out_list = [] |
| for _ in range(k): |
| val = torch_uniform(0, accumulate_weight_list[-1], generator) |
| active_id = 0 |
| for i, weight_val in enumerate(accumulate_weight_list): |
| active_id = i |
| if weight_val > val: |
| break |
| out_list.append(src_list[active_id]) |
|
|
| return out_list[0] if k == 1 else out_list |
|
|
