| from __future__ import annotations |
|
|
| import mlx.core as mx |
| import numpy as np |
|
|
|
|
| def patchify(latent): |
| """BCHW latent -> [B*H/2*W/2, C*2*2], channel then patch coordinates.""" |
| b, c, h, w = latent.shape |
| if h % 2 or w % 2: |
| raise ValueError("Latent height/width must be even") |
| return ( |
| latent.reshape(b, c, h // 2, 2, w // 2, 2) |
| .transpose(0, 2, 4, 1, 3, 5) |
| .reshape(-1, c * 4) |
| ) |
|
|
|
|
| def unpatchify(rows, h, w, channels=24): |
| return ( |
| rows.reshape(1, h // 2, w // 2, channels, 2, 2) |
| .transpose(0, 3, 1, 4, 2, 5) |
| .reshape(1, channels, h, w) |
| ) |
|
|
|
|
| def image_positions(text_len, h, w, references=()): |
| """H3 ref2va coordinates, generated in float64 on CPU then cast once. |
| |
| Padding is omitted: upstream packs it as a separate attention document. |
| Real rows cannot attend to it, so it has no effect on these outputs. |
| """ |
| text = np.zeros((text_len, 3), np.float64) |
| text[:, 0] = np.arange(text_len) |
| result = [text] |
| for i, (hh, ww) in enumerate((*references, (h, w))): |
| area = np.sqrt(hh * ww) |
| axes = [] |
| for dim in (hh, ww): |
| ratio = dim / area |
| left = (1 - ratio) / 2 |
| axes.append(np.linspace(left, left + ratio, dim // 2, endpoint=False) * 32) |
| yy, xx = np.meshgrid(*axes, indexing="ij") |
| coords = np.stack((np.full_like(yy, text_len + i), yy, xx), axis=-1).reshape( |
| -1, 3 |
| ) |
| result.append(coords) |
| return mx.array(np.concatenate(result).astype(np.float32)) |
|
|
|
|
| def rope_cache(positions, inv_freq, dtype=mx.bfloat16): |
| half = (positions[..., None] * inv_freq).reshape(positions.shape[0], -1) |
| angles = mx.concatenate((half, half), axis=-1) |
| return mx.cos(angles).astype(dtype)[None, :, None], mx.sin(angles).astype(dtype)[ |
| None, :, None |
| ] |
|
|
