(
vector: Vector,
dim: int,
device: torch.device,
)
| 16 | |
| 17 | |
| 18 | def sanitize_vector( |
| 19 | vector: Vector, |
| 20 | dim: int, |
| 21 | device: torch.device, |
| 22 | ) -> Float[Tensor, "*#batch dim"]: |
| 23 | if isinstance(vector, Tensor): |
| 24 | vector = vector.type(torch.float32).to(device) |
| 25 | else: |
| 26 | vector = torch.tensor(vector, dtype=torch.float32, device=device) |
| 27 | while vector.ndim < 2: |
| 28 | vector = vector[None] |
| 29 | if vector.shape[-1] == 1: |
| 30 | vector = repeat(vector, "... () -> ... c", c=dim) |
| 31 | assert vector.shape[-1] == dim |
| 32 | assert vector.ndim == 2 |
| 33 | return vector |
| 34 | |
| 35 | |
| 36 | Scalar = Union[ |
no test coverage detected