File size: 1,856 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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
]
|