File size: 4,096 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 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import pytest
mx = pytest.importorskip("mlx.core")
nn = pytest.importorskip("mlx.nn")
head_module = pytest.importorskip("gemma_rlcd.decision_head")
backend_module = pytest.importorskip("gemma_rlcd.head_backend")
HeadConfig = head_module.HeadConfig
ParallelDecisionHead = head_module.ParallelDecisionHead
grouped_cross_entropy = head_module.grouped_cross_entropy
DecisionHeadBackend = backend_module.DecisionHeadBackend
def fixture():
mx.random.seed(7)
model = ParallelDecisionHead(HeadConfig(input_dims=16, dims=8, heads=2))
model.eval()
state = mx.random.normal((1, 7, 16))
queries = mx.random.normal((5, 4, 16))
mask = mx.array([[True, True, False, False]] * 5)
fields = mx.array([0, 0, 1, 1, 1])
return model, state, queries, mask, fields
def test_all_fields_equal_independent_head_calls():
model, state, queries, mask, fields = fixture()
together = model(state, queries, mask, fields, 2)
separate = mx.concatenate(
[
model(state, queries[:2], mask[:2], mx.zeros(2, dtype=mx.int32), 1),
model(state, queries[2:], mask[2:], mx.zeros(3, dtype=mx.int32), 1),
]
)
assert bool(mx.allclose(together, separate, atol=1e-5).item())
def test_candidate_and_field_permutations_preserve_scores():
model, state, queries, mask, fields = fixture()
original = model(state, queries, mask, fields, 2)
permutation = mx.array([4, 2, 3, 1, 0])
reordered = model(state, queries[permutation], mask[permutation], mx.array([0, 0, 0, 1, 1]), 2)
assert bool(mx.allclose(reordered, original[permutation], atol=1e-5).item())
def test_padding_values_and_other_fields_cannot_change_a_field():
model, state, queries, mask, fields = fixture()
original = model(state, queries, mask, fields, 2)
altered = mx.where(mask[:, :, None], queries, queries + 1000)
padded = model(state, altered, mask, fields, 2)
assert bool(mx.allclose(original, padded, atol=1e-5).item())
altered = queries.at[2:].add(100)
changed = model(state, altered, mask, fields, 2)
assert bool(mx.allclose(original[:2], changed[:2], atol=1e-5).item())
def test_head_has_finite_nonzero_training_gradients():
model, state, queries, mask, fields = fixture()
def loss(head):
return grouped_cross_entropy(head(state, queries, mask, fields, 2), (0, 2, 5), (1, 2))
from mlx.utils import tree_flatten
value, gradients = nn.value_and_grad(model, loss)(model)
arrays = [array for _, array in tree_flatten(gradients)]
mx.eval(value, arrays)
assert bool(mx.isfinite(value).item())
assert all(bool(mx.all(mx.isfinite(array)).item()) for array in arrays)
assert sum(float(mx.sum(mx.abs(array)).item()) for array in arrays) > 0
def test_untrained_backend_cannot_return_decisions():
backend = DecisionHeadBackend.__new__(DecisionHeadBackend)
backend.training_metadata = None
with pytest.raises(RuntimeError, match="Untrained decision head"):
backend.score_batch(None, [])
def test_grouped_loss_rejects_invalid_target():
with pytest.raises(ValueError, match="Target index"):
grouped_cross_entropy(mx.zeros(5), (0, 2, 5), (2, 0))
def test_head_checkpoint_round_trip_preserves_candidate_scores(tmp_path):
model, state, queries, mask, fields = fixture()
original = model(state, queries, mask, fields, 2)
model.save_weights(str(tmp_path / "head.safetensors"))
restored = ParallelDecisionHead(model.config)
restored.load_weights(str(tmp_path / "head.safetensors"), strict=True)
restored.eval()
actual = restored(state, queries, mask, fields, 2)
assert bool(mx.array_equal(original, actual).item())
@pytest.mark.parametrize("updates", [None, 0, -1, True, 1.5, "600"])
def test_checkpoint_rejects_invalid_training_update_count(tmp_path, updates):
backend = DecisionHeadBackend.__new__(DecisionHeadBackend)
with pytest.raises(ValueError, match="positive integer training updates"):
backend.save(tmp_path / "invalid-head", {"updates": updates})
assert not (tmp_path / "invalid-head").exists()
|