Add GLM 5.3 Flash support

src/engine/metal/gpu.rs

@@ -5,6 +5,42 @@ pub(super) struct GpuTensor {
     _private: [u8; 0],
 }
 
+#[derive(Clone, Copy, Default)]
+#[repr(C)]
+pub(super) struct Glm53VisionLayerWeights {
+    pub(super) norm1: u64,
+    pub(super) qkv_weight: u64,
+    pub(super) qkv_bias: u64,
+    pub(super) q_norm: u64,
+    pub(super) k_norm: u64,
+    pub(super) attn_proj_weight: u64,
+    pub(super) attn_proj_bias: u64,
+    pub(super) norm2: u64,
+    pub(super) gate_weight: u64,
+    pub(super) gate_bias: u64,
+    pub(super) up_weight: u64,
+    pub(super) up_bias: u64,
+    pub(super) down_weight: u64,
+    pub(super) down_bias: u64,
+}
+
+#[derive(Clone, Copy, Default)]
+#[repr(C)]
+pub(super) struct Glm53VisionWeights {
+    pub(super) patch_weight: u64,
+    pub(super) patch_bias: u64,
+    pub(super) post_norm: u64,
+    pub(super) downsample_weight: u64,
+    pub(super) downsample_bias: u64,
+    pub(super) merger_proj: u64,
+    pub(super) merger_norm: u64,
+    pub(super) merger_norm_bias: u64,
+    pub(super) merger_gate: u64,
+    pub(super) merger_up: u64,
+    pub(super) merger_down: u64,
+    pub(super) layer: [Glm53VisionLayerWeights; 24],
+}
+
 #[repr(C)]
 pub(super) struct StreamExpertTable {
     pub(super) model_map: *const c_void,
@@ -93,6 +129,12 @@ unsafe extern "C" {
         count: u32,
     ) -> i32;
     pub(super) fn ds4_gpu_flush_commands() -> i32;
+    pub(super) fn ds4_gpu_flush_encoder() -> i32;
+    pub(super) fn ds4_gpu_argmax_tensor(
+        out: *mut GpuTensor,
+        logits: *const GpuTensor,
+        vocab: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32;
     pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32;
     #[cfg(test)]
@@ -227,6 +269,65 @@ unsafe extern "C" {
         embd: u32,
         hc: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_repeat_hc_rows_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        rows: u32,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_embedding_bf16(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        tokens: *const GpuTensor,
+        rows: u32,
+        embd: u32,
+        vocab: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_vision_encode(
+        out: *mut f32,
+        patches: *const f32,
+        grid_h: u32,
+        grid_w: u32,
+        map: *const c_void,
+        size: u64,
+        weights: *const Glm53VisionWeights,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_scatter_image_hc(
+        hc: *mut GpuTensor,
+        image: *const GpuTensor,
+        dst_row: u32,
+        image_row: u32,
+        rows: u32,
+        total_rows: u32,
+        embd: u32,
+        hc_count: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_matmul_bf16(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u32,
+        output: u32,
+        x: *const GpuTensor,
+        rows: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_matmul_bf16_qkv(
+        q: *mut GpuTensor,
+        k: *mut GpuTensor,
+        v: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        q_weight: u64,
+        k_weight: u64,
+        v_weight: u64,
+        input: u32,
+        output: u32,
+        x: *const GpuTensor,
+    ) -> i32;
     pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
         out: *mut GpuTensor,
         map: *const c_void,
@@ -417,6 +518,35 @@ unsafe extern "C" {
         beta_slow: f32,
         cache_f16: bool,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm53_indexer_pool_update_tensor(
+        cache: *mut GpuTensor,
+        tail_k: *mut GpuTensor,
+        tail_gate: *mut GpuTensor,
+        raw_k: *const GpuTensor,
+        gate: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        norm_weight: u64,
+        norm_bias: u64,
+        ape: u64,
+        pos: u32,
+        rows: u32,
+        cache_cap: u32,
+        head_dim: u32,
+        pool_size: u32,
+        eps: f32,
+        cache_f16: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_expand_pool_selection_tensor(
+        selected: *mut GpuTensor,
+        pools: *const GpuTensor,
+        rows: u32,
+        pos: u32,
+        selected_pools: u32,
+        top_k: u32,
+        pool_size: u32,
+        width: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_fill_selected_range_tensor(
         selected: *mut GpuTensor,
         count: u32,
@@ -467,6 +597,66 @@ unsafe extern "C" {
         scale: f32,
         cache_f16: bool,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm53_indexer_scores_batch_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        cache: *const GpuTensor,
+        visible: u32,
+        rows: u32,
+        pos: u32,
+        pool_size: u32,
+        heads: u32,
+        head_dim: u32,
+        scale: f32,
+        cache_f16: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_kda_decode(
+        out: *mut GpuTensor,
+        conv: *mut GpuTensor,
+        recurrent: *mut GpuTensor,
+        q: *const GpuTensor,
+        k: *const GpuTensor,
+        v: *const GpuTensor,
+        gate: *const GpuTensor,
+        beta: *const GpuTensor,
+        output_gate: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        q_conv: u64,
+        k_conv: u64,
+        v_conv: u64,
+        a_log: u64,
+        dt_bias: u64,
+        output_norm: u64,
+        heads: u32,
+        rows: u32,
+        gate_lower_bound: f32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm53_kda_prefill(
+        out: *mut GpuTensor,
+        conv: *mut GpuTensor,
+        recurrent: *mut GpuTensor,
+        q: *mut GpuTensor,
+        k: *mut GpuTensor,
+        v: *mut GpuTensor,
+        gate: *mut GpuTensor,
+        beta: *const GpuTensor,
+        output_gate: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        q_conv: u64,
+        k_conv: u64,
+        v_conv: u64,
+        a_log: u64,
+        dt_bias: u64,
+        output_norm: u64,
+        heads: u32,
+        rows: u32,
+        gate_lower_bound: f32,
+        eps: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
         out: *mut GpuTensor,
         q: *const GpuTensor,
@@ -554,6 +744,19 @@ unsafe extern "C" {
         beta_fast: f32,
         beta_slow: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
+        out: *mut GpuTensor,
+        qk_low: *const GpuTensor,
+        kv_cache: *const GpuTensor,
+        q_row0: u32,
+        rows: u32,
+        selected: u32,
+        cache_cap: u32,
+        cache_f16: bool,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
         out: *mut GpuTensor,
         q: *const GpuTensor,
@@ -777,6 +980,20 @@ unsafe extern "C" {
         eps: f32,
         norm_eps: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_hc_split_weighted_sum_tensor(
+        out: *mut GpuTensor,
+        split: *mut GpuTensor,
+        mix: *const GpuTensor,
+        residual: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        scale: u64,
+        base: u64,
+        embd: u32,
+        hc: u32,
+        iterations: u32,
+        eps: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_available() -> i32;
     pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_tensor(
         out: *mut GpuTensor,
@@ -1359,6 +1576,15 @@ unsafe extern "C" {
         embd: u32,
         hc: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_hc_expand_tensor(
+        out: *mut GpuTensor,
+        block: *const GpuTensor,
+        residual: *const GpuTensor,
+        post: *const GpuTensor,
+        combine: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_hc_expand_add_split_tensor(
         out: *mut GpuTensor,
         block: *const GpuTensor,
@@ -1463,6 +1689,44 @@ unsafe extern "C" {
         x: *const GpuTensor,
         clamp: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_shared_mid_swiglu_q8_0_tensor(
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        clamp: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
+        gate: *mut GpuTensor,
+        up: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        clamp: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
+        gate: *mut GpuTensor,
+        up: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+        clamp: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_router_shared_gate_up_q8_0_tensor(
         router_logits: *mut GpuTensor,
         gate: *mut GpuTensor,
@@ -1598,6 +1862,21 @@ impl Context {
                 return Err(error);
             }
         }
+        if let Some(vision) = &model.vision {
+            let mapped = unsafe {
+                ds4_gpu_set_model_map_range(
+                    vision.map_ptr().cast(),
+                    vision.len(),
+                    vision.data_offset(),
+                    vision.len() - vision.data_offset(),
+                    vision.max_tensor_bytes(),
+                )
+            };
+            if let Err(error) = check(mapped, "vision-model mapping") {
+                unsafe { ds4_gpu_cleanup() };
+                return Err(error);
+            }
+        }
         unsafe { ds4_gpu_set_quality(quality) };
         let model_file = File::open(model.main.path()).map_err(|error| {
             unsafe { ds4_gpu_cleanup() };

src/engine/metal/vision.rs

@@ -0,0 +1,311 @@
+use super::gpu::{Glm53VisionLayerWeights, Glm53VisionWeights, ds4_gpu_glm53_vision_encode};
+use super::{Gguf, call};
+use image::{DynamicImage, ImageDecoder, ImageReader};
+use std::io::Cursor;
+
+const EMBEDDING: usize = 4096;
+const PATCH: u32 = 14;
+const MERGE: u32 = 2;
+const MIN_TOKENS: u32 = 16;
+const MAX_TOKENS: u32 = 8000;
+
+pub(super) struct VisionEncoder {
+    weights: Glm53VisionWeights,
+}
+
+pub(in crate::engine) struct VisionEmbedding {
+    pub(in crate::engine) values: Vec<f32>,
+    pub(in crate::engine) tokens: u32,
+    pub(in crate::engine) width: u32,
+    pub(in crate::engine) height: u32,
+    pub(in crate::engine) content_width: u32,
+    pub(in crate::engine) content_height: u32,
+}
+
+impl VisionEncoder {
+    pub(super) fn bind(model: &Gguf) -> Result<Self, String> {
+        let offset = |name: &str| model.tensor(name).map(|tensor| tensor.offset);
+        let mut weights = Glm53VisionWeights {
+            patch_weight: offset("model.visual.patch_embed.proj.weight")?,
+            patch_bias: offset("model.visual.patch_embed.proj.bias")?,
+            post_norm: offset("model.visual.post_layernorm.weight")?,
+            downsample_weight: offset("model.visual.downsample.weight")?,
+            downsample_bias: offset("model.visual.downsample.bias")?,
+            merger_proj: offset("model.visual.merger.proj.weight")?,
+            merger_norm: offset("model.visual.merger.post_projection_norm.weight")?,
+            merger_norm_bias: offset("model.visual.merger.post_projection_norm.bias")?,
+            merger_gate: offset("model.visual.merger.gate_proj.weight")?,
+            merger_up: offset("model.visual.merger.up_proj.weight")?,
+            merger_down: offset("model.visual.merger.down_proj.weight")?,
+            ..Glm53VisionWeights::default()
+        };
+        for (layer, target) in weights.layer.iter_mut().enumerate() {
+            let name = |suffix: &str| format!("model.visual.blocks.{layer}.{suffix}");
+            *target = Glm53VisionLayerWeights {
+                norm1: offset(&name("norm1.weight"))?,
+                qkv_weight: offset(&name("attn.qkv.weight"))?,
+                qkv_bias: offset(&name("attn.qkv.bias"))?,
+                q_norm: offset(&name("attn.q_norm.weight"))?,
+                k_norm: offset(&name("attn.k_norm.weight"))?,
+                attn_proj_weight: offset(&name("attn.proj.weight"))?,
+                attn_proj_bias: offset(&name("attn.proj.bias"))?,
+                norm2: offset(&name("norm2.weight"))?,
+                gate_weight: offset(&name("mlp.gate_proj.weight"))?,
+                gate_bias: offset(&name("mlp.gate_proj.bias"))?,
+                up_weight: offset(&name("mlp.up_proj.weight"))?,
+                up_bias: offset(&name("mlp.up_proj.bias"))?,
+                down_weight: offset(&name("mlp.down_proj.weight"))?,
+                down_bias: offset(&name("mlp.down_proj.bias"))?,
+            };
+        }
+        Ok(Self { weights })
+    }
+
+    pub(super) fn encode(&self, model: &Gguf, encoded: &[u8]) -> Result<VisionEmbedding, String> {
+        if encoded.is_empty() || encoded.len() > 64 * 1024 * 1024 {
+            return Err("image is empty or exceeds the 64 MiB encoded limit".into());
+        }
+        let reader = ImageReader::new(Cursor::new(encoded))
+            .with_guessed_format()
+            .map_err(|error| error.to_string())?;
+        let mut decoder = reader.into_decoder().map_err(|error| error.to_string())?;
+        let orientation = decoder.orientation().map_err(|error| error.to_string())?;
+        let mut image = DynamicImage::from_decoder(decoder).map_err(|error| error.to_string())?;
+        image.apply_orientation(orientation);
+        let rgb = image.into_rgb8();
+        let (width, height) = rgb.dimensions();
+        if width == 0
+            || height == 0
+            || width > 16_384
+            || height > 16_384
+            || u64::from(width) * u64::from(height) > 64 * 1024 * 1024
+        {
+            return Err("image dimensions exceed the GLM 5.3 vision limits".into());
+        }
+        let patches = preprocess(rgb.as_raw(), width, height)?;
+        let tokens = patches.grid_height * patches.grid_width / 4;
+        let mut values = vec![0.0_f32; tokens as usize * EMBEDDING];
+        call(
+            unsafe {
+                ds4_gpu_glm53_vision_encode(
+                    values.as_mut_ptr(),
+                    patches.values.as_ptr(),
+                    patches.grid_height,
+                    patches.grid_width,
+                    model.map_ptr().cast(),
+                    model.len(),
+                    &self.weights,
+                )
+            },
+            "encoding a GLM 5.3 image",
+        )?;
+        Ok(VisionEmbedding {
+            values,
+            tokens,
+            width,
+            height,
+            content_width: patches.content_width,
+            content_height: patches.content_height,
+        })
+    }
+}
+
+struct Patches {
+    values: Vec<f32>,
+    content_width: u32,
+    content_height: u32,
+    grid_width: u32,
+    grid_height: u32,
+}
+
+fn preprocess(rgb: &[u8], width: u32, height: u32) -> Result<Patches, String> {
+    const MEAN: [f32; 3] = [0.48145466, 0.4578275, 0.40821073];
+    const STDDEV: [f32; 3] = [0.26862954, 0.261_302_6, 0.275_777_1];
+    let (target_height, target_width) = smart_resize(height, width)?;
+    let mut scale = (target_height as f64 / height as f64).min(target_width as f64 / width as f64);
+    if 2 * u64::from(height) * u64::from(width) >= 2 * 28 * 28 * MIN_TOKENS as u64 && scale > 1.0 {
+        scale = 1.0;
+    }
+    let content_height = ((height as f64 * scale).floor() as u32).clamp(1, target_height);
+    let content_width = ((width as f64 * scale).floor() as u32).clamp(1, target_width);
+    let mut canvas = vec![0.0_f32; target_height as usize * target_width as usize * 3];
+    if content_width == width && content_height == height {
+        for y in 0..height {
+            for x in 0..width {
+                let source = (y as usize * width as usize + x as usize) * 3;
+                let target = (y as usize * target_width as usize + x as usize) * 3;
+                for channel in 0..3 {
+                    canvas[target + channel] = f32::from(rgb[source + channel]);
+                }
+            }
+        }
+    } else {
+        resize_bicubic(
+            rgb,
+            width,
+            height,
+            &mut canvas,
+            content_width,
+            content_height,
+            target_width,
+        );
+    }
+    for y in 0..target_height {
+        for x in 0..target_width {
+            let pixel = (y as usize * target_width as usize + x as usize) * 3;
+            for channel in 0..3 {
+                let value = if x < content_width && y < content_height {
+                    canvas[pixel + channel]
+                } else {
+                    0.0
+                };
+                canvas[pixel + channel] = (value / 255.0 - MEAN[channel]) / STDDEV[channel];
+            }
+        }
+    }
+
+    let grid_height = target_height / PATCH;
+    let grid_width = target_width / PATCH;
+    let patch_values = grid_height as usize * grid_width as usize * 3 * 2 * 14 * 14;
+    let mut values = Vec::with_capacity(patch_values);
+    for block_y in 0..grid_height / MERGE {
+        for block_x in 0..grid_width / MERGE {
+            for merge_y in 0..MERGE {
+                for merge_x in 0..MERGE {
+                    let patch_y = block_y * MERGE + merge_y;
+                    let patch_x = block_x * MERGE + merge_x;
+                    for channel in 0..3 {
+                        for _ in 0..2 {
+                            for y in 0..PATCH {
+                                for x in 0..PATCH {
+                                    let pixel = ((patch_y * PATCH + y) as usize
+                                        * target_width as usize
+                                        + (patch_x * PATCH + x) as usize)
+                                        * 3;
+                                    values.push(canvas[pixel + channel]);
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    if values.len() != patch_values {
+        return Err("internal GLM 5.3 vision patch layout mismatch".into());
+    }
+    Ok(Patches {
+        values,
+        content_width,
+        content_height,
+        grid_width,
+        grid_height,
+    })
+}
+
+fn smart_resize(height: u32, width: u32) -> Result<(u32, u32), String> {
+    let factor = 28_u32;
+    let align = |value: u32| value.div_ceil(factor) * factor;
+    let pixels_per_token = 2_u64 * factor as u64 * factor as u64;
+    let min_pixels = MIN_TOKENS as u64 * pixels_per_token;
+    let max_pixels = MAX_TOKENS as u64 * pixels_per_token;
+    let mut aligned_height = align(height);
+    let mut aligned_width = align(width);
+    let mut budget = 2_u64 * aligned_height as u64 * aligned_width as u64;
+    if budget < min_pixels {
+        let scale = (min_pixels as f64 / (2.0 * height as f64 * width as f64)).sqrt();
+        aligned_height = align((height as f64 * scale).ceil() as u32);
+        aligned_width = align((width as f64 * scale).ceil() as u32);
+        budget = 2_u64 * aligned_height as u64 * aligned_width as u64;
+    }
+    if budget > max_pixels {
+        let (mut low, mut high) = (1_u32, height);
+        aligned_height = factor;
+        aligned_width = factor;
+        while low <= high {
+            let content_height = low + (high - low) / 2;
+            let content_width =
+                ((width as f64 * content_height as f64 / height as f64).floor() as u32).max(1);
+            let candidate_height = align(content_height);
+            let candidate_width = align(content_width);
+            if 2_u64 * candidate_height as u64 * candidate_width as u64 <= max_pixels {
+                aligned_height = candidate_height;
+                aligned_width = candidate_width;
+                low = content_height + 1;
+            } else {
+                high = content_height - 1;
+            }
+        }
+    }
+    Ok((aligned_height, aligned_width))
+}
+
+#[allow(clippy::too_many_arguments)]
+fn resize_bicubic(
+    source: &[u8],
+    source_width: u32,
+    source_height: u32,
+    target: &mut [f32],
+    target_width: u32,
+    target_height: u32,
+    target_stride: u32,
+) {
+    let scale_x = source_width as f64 / target_width as f64;
+    let scale_y = source_height as f64 / target_height as f64;
+    let filter_x = if scale_x >= 1.0 { 1.0 / scale_x } else { 1.0 };
+    let filter_y = if scale_y >= 1.0 { 1.0 / scale_y } else { 1.0 };
+    let support_x = if scale_x >= 1.0 { 2.0 * scale_x } else { 2.0 };
+    let support_y = if scale_y >= 1.0 { 2.0 * scale_y } else { 2.0 };
+    for dy in 0..target_height {
+        let center_y = scale_y * (dy as f64 + 0.5);
+        let y0 = (center_y - support_y + 0.5).max(0.0) as u32;
+        let y1 = (center_y + support_y + 0.5).min(source_height as f64) as u32;
+        for dx in 0..target_width {
+            let center_x = scale_x * (dx as f64 + 0.5);
+            let x0 = (center_x - support_x + 0.5).max(0.0) as u32;
+            let x1 = (center_x + support_x + 0.5).min(source_width as f64) as u32;
+            let mut sum = [0.0; 3];
+            let mut weight_sum = 0.0;
+            for iy in y0..y1 {
+                let wy = cubic((iy as f64 + 0.5 - center_y) * filter_y);
+                for ix in x0..x1 {
+                    let weight = wy * cubic((ix as f64 + 0.5 - center_x) * filter_x);
+                    let pixel = (iy as usize * source_width as usize + ix as usize) * 3;
+                    for channel in 0..3 {
+                        sum[channel] += source[pixel + channel] as f64 * weight;
+                    }
+                    weight_sum += weight;
+                }
+            }
+            let pixel = (dy as usize * target_stride as usize + dx as usize) * 3;
+            for channel in 0..3 {
+                target[pixel + channel] =
+                    (sum[channel] / weight_sum).round().clamp(0.0, 255.0) as f32;
+            }
+        }
+    }
+}
+
+fn cubic(mut x: f64) -> f64 {
+    const A: f64 = -0.5;
+    x = x.abs();
+    if x < 1.0 {
+        ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0
+    } else if x < 2.0 {
+        ((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A
+    } else {
+        0.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::smart_resize;
+
+    #[test]
+    fn glm53_resize_matches_reference_token_grids() {
+        assert_eq!(smart_resize(28, 28), Ok((112, 112)));
+        assert_eq!(smart_resize(1024, 1024), Ok((1036, 1036)));
+        assert_eq!(smart_resize(1080, 1920), Ok((1092, 1932)));
+    }
+}