| import torch |
| from typing import Tuple, Callable |
| def hacer_nada(x: torch.Tensor, modo: str = None): |
| return x |
| def brujeria_mps(entrada, dim, indice): |
| if entrada.shape[-1] == 1: |
| return torch.gather(entrada.unsqueeze(-1), dim - 1 if dim < 0 else dim, indice.unsqueeze(-1)).squeeze(-1) |
| else: |
| return torch.gather(entrada, dim, indice) |
| def emparejamiento_suave_aleatorio_2d( |
| metrica: torch.Tensor, |
| ancho: int, |
| alto: int, |
| paso_x: int, |
| paso_y: int, |
| radio: int, |
| sin_aleatoriedad: bool = False, |
| generador: torch.Generator = None |
| ) -> Tuple[Callable, Callable]: |
| lote, num_nodos, _ = metrica.shape |
| if radio <= 0: |
| return hacer_nada, hacer_nada |
| recopilar = brujeria_mps if metrica.device.type == "mps" else torch.gather |
| with torch.no_grad(): |
| alto_paso_y, ancho_paso_x = alto // paso_y, ancho // paso_x |
| if sin_aleatoriedad: |
| indice_aleatorio = torch.zeros(alto_paso_y, ancho_paso_x, 1, device=metrica.device, dtype=torch.int64) |
| else: |
| indice_aleatorio = torch.randint(paso_y * paso_x, size=(alto_paso_y, ancho_paso_x, 1), device=generador.device, generator=generador).to(metrica.device) |
| vista_buffer_indice = torch.zeros(alto_paso_y, ancho_paso_x, paso_y * paso_x, device=metrica.device, dtype=torch.int64) |
| vista_buffer_indice.scatter_(dim=2, index=indice_aleatorio, src=-torch.ones_like(indice_aleatorio, dtype=indice_aleatorio.dtype)) |
| vista_buffer_indice = vista_buffer_indice.view(alto_paso_y, ancho_paso_x, paso_y, paso_x).transpose(1, 2).reshape(alto_paso_y * paso_y, ancho_paso_x * paso_x) |
| if (alto_paso_y * paso_y) < alto or (ancho_paso_x * paso_x) < ancho: |
| buffer_indice = torch.zeros(alto, ancho, device=metrica.device, dtype=torch.int64) |
| buffer_indice[:(alto_paso_y * paso_y), :(ancho_paso_x * paso_x)] = vista_buffer_indice |
| else: |
| buffer_indice = vista_buffer_indice |
| indice_aleatorio = buffer_indice.reshape(1, -1, 1).argsort(dim=1) |
| del buffer_indice, vista_buffer_indice |
| num_destino = alto_paso_y * ancho_paso_x |
| indices_a = indice_aleatorio[:, num_destino:, :] |
| indices_b = indice_aleatorio[:, :num_destino, :] |
| def dividir(x): |
| canales = x.shape[-1] |
| origen = recopilar(x, dim=1, index=indices_a.expand(lote, num_nodos - num_destino, canales)) |
| destino = recopilar(x, dim=1, index=indices_b.expand(lote, num_destino, canales)) |
| return origen, destino |
| metrica = metrica / metrica.norm(dim=-1, keepdim=True) |
| a, b = dividir(metrica) |
| puntuaciones = a @ b.transpose(-1, -2) |
| radio = min(a.shape[1], radio) |
| nodo_max, nodo_indice = puntuaciones.max(dim=-1) |
| indice_borde = nodo_max.argsort(dim=-1, descending=True)[..., None] |
| indice_no_emparejado = indice_borde[..., radio:, :] |
| indice_origen = indice_borde[..., :radio, :] |
| indice_destino = recopilar(nodo_indice[..., None], dim=-2, index=indice_origen) |
| def fusionar(x: torch.Tensor, modo="mean") -> torch.Tensor: |
| origen, destino = dividir(x) |
| n, t1, c = origen.shape |
| no_emparejado = recopilar(origen, dim=-2, index=indice_no_emparejado.expand(n, t1 - radio, c)) |
| origen = recopilar(origen, dim=-2, index=indice_origen.expand(n, radio, c)) |
| destino = destino.scatter_reduce(-2, indice_destino.expand(n, radio, c), origen, reduce=modo) |
| return torch.cat([no_emparejado, destino], dim=1) |
| def desfusionar(x: torch.Tensor) -> torch.Tensor: |
| longitud_no_emparejado = indice_no_emparejado.shape[1] |
| no_emparejado, destino = x[..., :longitud_no_emparejado, :], x[..., longitud_no_emparejado:, :] |
| _, _, c = no_emparejado.shape |
| origen = recopilar(destino, dim=-2, index=indice_destino.expand(lote, radio, c)) |
| salida = torch.zeros(lote, num_nodos, c, device=x.device, dtype=x.dtype) |
| salida.scatter_(dim=-2, index=indices_b.expand(lote, num_destino, c), src=destino) |
| salida.scatter_(dim=-2, index=recopilar(indices_a.expand(lote, indices_a.shape[1], 1), dim=1, index=indice_no_emparejado).expand(lote, longitud_no_emparejado, c), src=no_emparejado) |
| salida.scatter_(dim=-2, index=recopilar(indices_a.expand(lote, indices_a.shape[1], 1), dim=1, index=indice_origen).expand(lote, radio, c), src=origen) |
| return salida |
| return fusionar, desfusionar |
