File size: 1,183 Bytes
53e24ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | import pytest
mx = pytest.importorskip("mlx.core")
vision = pytest.importorskip("gemma_rlcd.vision")
config_module = pytest.importorskip("mlx_vlm.models.gemma4.config")
@pytest.mark.parametrize("dtype", [mx.float16, mx.float32, mx.bfloat16])
def test_indexed_positions_match_dense_lookup_including_padding(dtype):
config = config_module.VisionConfig(hidden_size=8, position_embedding_size=16)
original = vision.VisionPatchEmbedder(config)
original.position_embedding_table = mx.random.normal((2, 16, 8)).astype(dtype)
indexed = vision.IndexedVisionPatchEmbedder(original)
positions = mx.array([[[2, 3], [0, 0], [-1, -1], [16, 2], [4, 5]]])
padding = mx.array([[False, False, True, False, True]])
expected = original._position_embeddings(positions, padding)
actual = indexed._position_embeddings(positions, padding)
# The M5 dense float32 matmul can round table entries internally; gather
# returns their stored values directly. Half/bfloat paths match exactly.
if dtype == mx.float32:
assert bool(mx.allclose(actual, expected, atol=2e-3, rtol=2e-3).item())
else:
assert bool(mx.array_equal(actual, expected).item())
|