Draws stratified samples from a uniform distribution.
(shape, group=0, groups=1, dtype=None, device=None)
| 266 | |
| 267 | |
| 268 | def stratified_uniform(shape, group=0, groups=1, dtype=None, device=None): |
| 269 | """Draws stratified samples from a uniform distribution.""" |
| 270 | if groups <= 0: |
| 271 | raise ValueError(f"groups must be positive, got {groups}") |
| 272 | if group < 0 or group >= groups: |
| 273 | raise ValueError(f"group must be in [0, {groups})") |
| 274 | n = shape[-1] * groups |
| 275 | offsets = torch.arange(group, n, groups, dtype=dtype, device=device) |
| 276 | u = torch.rand(shape, dtype=dtype, device=device) |
| 277 | return (offsets + u) / n |
| 278 | |
| 279 | |
| 280 | stratified_settings = threading.local() |
no outgoing calls
no test coverage detected