Scatter the minimum value along the given dimension of `input` into `src` at the indices specified in `index`.
(size: int, dim: int, index: torch.LongTensor, src: torch.Tensor)
| 11 | |
| 12 | |
| 13 | def scatter_min(size: int, dim: int, index: torch.LongTensor, src: torch.Tensor) -> torch.return_types.min: |
| 14 | "Scatter the minimum value along the given dimension of `input` into `src` at the indices specified in `index`." |
| 15 | shape = src.shape[:dim] + (size,) + src.shape[dim + 1:] |
| 16 | minimum = torch.full(shape, float('inf'), dtype=src.dtype, device=src.device).scatter_reduce(dim=dim, index=index, src=src, reduce='amin', include_self=False) |
| 17 | minimum_where = torch.where(src == torch.gather(minimum, dim=dim, index=index)) |
| 18 | indices = torch.full(shape, -1, dtype=torch.long, device=src.device) |
| 19 | indices[(*minimum_where[:dim], index[minimum_where], *minimum_where[dim + 1:])] = minimum_where[dim] |
| 20 | return torch.return_types.min((minimum, indices)) |
| 21 | |
| 22 | |
| 23 | def split_batch_fwd(fn: Callable, chunk_size: int, *args, **kwargs): |
no outgoing calls
no test coverage detected