k3-node 1.0.0__py3-none-any.whl

k3_node/models/test_grover.py

@@ -0,0 +1,287 @@
+import os
+import pytest
+import numpy as np
+import keras
+from keras import ops
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.models.grover import (
+    GROVER,
+    GroverPReLU,
+    MPNEncoder,
+    MTBlock,
+    Readout,
+    GTransEncoder,
+    load_grover_weights,
+    download_grover_checkpoint,
+)
+
+
+def _make_dummy_batch(num_atoms=5, num_bonds=7, atom_fdim=151, bond_fdim=165, edge_fdim=None):
+    if edge_fdim is not None:
+        bond_fdim = edge_fdim
+    np.random.seed(42)
+    f_atoms = np.random.randn(num_atoms, atom_fdim).astype(np.float32)
+    f_atoms[0] = 0.0
+    f_bonds = np.random.randn(num_bonds, bond_fdim).astype(np.float32)
+    f_bonds[0] = 0.0
+
+    a2b = np.array([[0, 0], [2, 0], [1, 4], [3, 0], [0, 0]], dtype=np.int64)
+    b2a = np.array([0, 1, 2, 2, 3, 1, 3], dtype=np.int64)
+    b2revb = np.array([0, 2, 1, 4, 3, 6, 5], dtype=np.int64)
+    a_scope = np.array([[1, 2], [3, 2]], dtype=np.int64)
+    b_scope = np.array([[1, 3], [4, 3]], dtype=np.int64)
+    a2a = b2a[a2b]
+
+    return (
+        ops.convert_to_tensor(f_atoms),
+        ops.convert_to_tensor(f_bonds),
+        ops.convert_to_tensor(a2b),
+        ops.convert_to_tensor(b2a),
+        ops.convert_to_tensor(b2revb),
+        ops.convert_to_tensor(a_scope),
+        ops.convert_to_tensor(b_scope),
+        ops.convert_to_tensor(a2a),
+    )
+
+
+def test_grover_prelu():
+    layer = GroverPReLU(init_val=0.25)
+    x = ops.convert_to_tensor([-2.0, 0.0, 3.0], dtype="float32")
+    y = layer(x)
+    np.testing.assert_allclose(ops.convert_to_numpy(y), [-0.5, 0.0, 3.0], atol=1e-6)
+
+    # Config test
+    config = layer.get_config()
+    assert config["init_val"] == 0.25
+
+
+def test_mpn_encoder():
+    hidden_size = 32
+    depth = 3
+    num_atoms = 4
+    num_bonds = 4
+    atom_dim = 16
+    bond_dim = 20
+
+    # 1. Atom messages
+    encoder_atom = MPNEncoder(
+        hidden_size=hidden_size,
+        depth=depth,
+        atom_messages=True,
+        input_layer="fc",
+        input_dim=atom_dim,
+    )
+    init_msg = ops.convert_to_tensor(np.random.randn(num_atoms, atom_dim).astype(np.float32))
+    a2nei = ops.convert_to_tensor(np.array([[0, 0], [2, 0], [1, 3], [2, 0]]), dtype="int64")
+
+    out = encoder_atom(
+        init_messages=init_msg,
+        init_attached_features=None,
+        a2nei=a2nei,
+        a2attached=None,
+    )
+    assert ops.shape(out) == (num_atoms, hidden_size)
+
+    # 2. Bond messages (non-backtracking)
+    encoder_bond = MPNEncoder(
+        hidden_size=hidden_size,
+        depth=depth,
+        atom_messages=False,
+        input_layer="fc",
+        input_dim=bond_dim,
+    )
+    init_bond = ops.convert_to_tensor(np.random.randn(num_bonds, bond_dim).astype(np.float32))
+    a2b = ops.convert_to_tensor(np.array([[0], [2], [1], [3]]), dtype="int64")
+    b2a = ops.convert_to_tensor(np.array([0, 1, 2, 2]), dtype="int64")
+    b2revb = ops.convert_to_tensor(np.array([0, 2, 1, 3]), dtype="int64")
+
+    out_bond = encoder_bond(
+        init_messages=init_bond,
+        init_attached_features=None,
+        a2nei=a2b,
+        a2attached=None,
+        b2a=b2a,
+        b2revb=b2revb,
+    )
+    assert ops.shape(out_bond) == (num_bonds, hidden_size)
+
+
+def test_mt_block():
+    hidden_size = 32
+    num_attn_head = 2
+    depth = 3
+    num_atoms = 5
+    num_bonds = 7
+    atom_dim = 16
+    bond_dim = 20
+
+    block = MTBlock(
+        hidden_size=hidden_size,
+        input_dim=bond_dim,
+        num_attn_head=num_attn_head,
+        depth=depth,
+        atom_messages=False,
+    )
+
+    batch = _make_dummy_batch(num_atoms, num_bonds, atom_dim, bond_dim)
+    f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a = batch
+
+    out_atoms, out_bonds = block(f_atoms, f_bonds, a2b, b2a, b2revb, a2a)
+    assert ops.shape(out_bonds) == (num_bonds, hidden_size)
+
+
+def test_readout():
+    hidden_size = 32
+    emb = ops.convert_to_tensor(np.random.randn(8, hidden_size).astype(np.float32))
+    scope = ops.convert_to_tensor(np.array([[0, 3], [3, 4], [7, 1]]), dtype="int64")
+
+    # Mean Readout
+    readout_mean = Readout(rtype="mean", hidden_size=hidden_size)
+    out_mean = readout_mean(emb, scope)
+    assert ops.shape(out_mean) == (3, hidden_size)
+
+    # Self-attention Readout
+    readout_sa = Readout(rtype="self_attention", hidden_size=hidden_size, attn_hidden=16, attn_out=4)
+    out_sa = readout_sa(emb, scope)
+    assert ops.shape(out_sa) == (3, 4 * hidden_size)
+
+
+def test_gtrans_encoder():
+    hidden_size = 32
+    batch = _make_dummy_batch(num_atoms=5, num_bonds=7, atom_fdim=151, edge_fdim=165)
+    f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a = batch
+
+    # 'both'
+    encoder = GTransEncoder(
+        hidden_size=hidden_size,
+        edge_fdim=165,
+        node_fdim=151,
+        num_mt_block=1,
+        num_attn_head=2,
+        depth=3,
+        atom_emb_output="both",
+    )
+    res = encoder(f_atoms, f_bonds, a2b, b2a, b2revb, a2a)
+    assert "atom_from_atom" in res
+    assert "bond_from_bond" in res
+    assert ops.shape(res["atom_from_atom"]) == (5, hidden_size)
+    assert ops.shape(res["bond_from_bond"]) == (7, hidden_size)
+
+
+def test_grover_model_and_fingerprint():
+    hidden_size = 32
+    model = GROVER(
+        hidden_size=hidden_size,
+        edge_fdim=165,
+        node_fdim=151,
+        num_mt_block=1,
+        num_attn_head=2,
+        depth=3,
+        atom_emb_output="both",
+        readout_type="mean",
+    )
+
+    batch = _make_dummy_batch(num_atoms=5, num_bonds=7, atom_fdim=151, edge_fdim=165)
+    out = model(batch)
+    assert ops.shape(out["atom_from_atom"]) == (5, hidden_size)
+
+    # Fingerprint
+    fp = model.get_fingerprint(batch, fingerprint_source="both")
+    # 2 molecules in scope, 4 branches * 32 = 128
+    assert ops.shape(fp) == (2, 4 * hidden_size)
+
+
+def test_load_grover_weights_synthetic(tmp_path):
+    if torch is None:
+        pytest.skip("PyTorch is required for checkpoint loading test")
+    hidden_size = 16
+    num_attn_head = 2
+    depth = 2
+    node_fdim = 151
+    edge_fdim = 165
+
+    model = GROVER(
+        hidden_size=hidden_size,
+        edge_fdim=edge_fdim,
+        node_fdim=node_fdim,
+        num_mt_block=1,
+        num_attn_head=num_attn_head,
+        depth=depth,
+        atom_emb_output="both",
+    )
+    model.build(None)
+
+    # Create dummy torch state dict
+    state_dict = {}
+    for block_name, dim in [("edge_blocks", edge_fdim), ("node_blocks", node_fdim)]:
+        prefix = f"grover.encoders.{block_name}.0"
+        state_dict[f"{prefix}.W_i.weight"] = torch.randn(hidden_size, dim)
+        state_dict[f"{prefix}.act_func.weight"] = torch.tensor([0.25])
+        state_dict[f"{prefix}.layernorm.weight"] = torch.ones(hidden_size)
+        state_dict[f"{prefix}.layernorm.bias"] = torch.zeros(hidden_size)
+        state_dict[f"{prefix}.sublayer.norm.weight"] = torch.ones(hidden_size)
+        state_dict[f"{prefix}.sublayer.norm.bias"] = torch.zeros(hidden_size)
+        state_dict[f"{prefix}.W_o.weight"] = torch.randn(hidden_size, hidden_size * num_attn_head)
+        for i in range(3):
+            state_dict[f"{prefix}.attn.linear_layers.{i}.weight"] = torch.randn(hidden_size, hidden_size)
+            state_dict[f"{prefix}.attn.linear_layers.{i}.bias"] = torch.zeros(hidden_size)
+        state_dict[f"{prefix}.attn.output_linear.weight"] = torch.randn(hidden_size, hidden_size)
+        for hi in range(num_attn_head):
+            for mpn in ["mpn_q", "mpn_k", "mpn_v"]:
+                state_dict[f"{prefix}.heads.{hi}.{mpn}.W_h.weight"] = torch.randn(hidden_size, hidden_size)
+                state_dict[f"{prefix}.heads.{hi}.{mpn}.act_func.weight"] = torch.tensor([0.25])
+
+    for ffn_key, in_dim in [
+        ("ffn_atom_from_atom", hidden_size + node_fdim),
+        ("ffn_atom_from_bond", hidden_size + node_fdim),
+        ("ffn_bond_from_atom", hidden_size + edge_fdim),
+        ("ffn_bond_from_bond", hidden_size + edge_fdim),
+    ]:
+        prefix = f"grover.encoders.{ffn_key}"
+        state_dict[f"{prefix}.W_1.weight"] = torch.randn(hidden_size * 4, in_dim)
+        state_dict[f"{prefix}.W_1.bias"] = torch.zeros(hidden_size * 4)
+        state_dict[f"{prefix}.W_2.weight"] = torch.randn(hidden_size, hidden_size * 4)
+        state_dict[f"{prefix}.W_2.bias"] = torch.zeros(hidden_size)
+        state_dict[f"{prefix}.act_func.weight"] = torch.tensor([0.25])
+
+    for sub in ["atom_from_atom", "atom_from_bond", "bond_from_atom", "bond_from_bond"]:
+        state_dict[f"grover.encoders.{sub}_sublayer.norm.weight"] = torch.ones(hidden_size)
+        state_dict[f"grover.encoders.{sub}_sublayer.norm.bias"] = torch.zeros(hidden_size)
+
+    state_dict["grover.encoders.act_func_node.weight"] = torch.tensor([0.25])
+    state_dict["grover.encoders.act_func_edge.weight"] = torch.tensor([0.25])
+
+    save_path = str(tmp_path / "mock_grover.pt")
+    torch.save({"state_dict": state_dict}, save_path)
+
+    load_grover_weights(model, save_path)
+    # verify forward pass succeeds
+    batch = _make_dummy_batch(num_atoms=5, num_bonds=7, atom_fdim=node_fdim, edge_fdim=edge_fdim)
+    out = model(batch)
+    assert ops.shape(out["atom_from_atom"]) == (5, hidden_size)
+
+
+def test_grover_official_pretrained_load():
+    ckpt_path = "/tmp/grover_download_test/grover_base.pt"
+    if not os.path.exists(ckpt_path):
+        pytest.skip("Official checkpoint not present at /tmp/grover_download_test/grover_base.pt")
+
+    model = GROVER(
+        hidden_size=800,
+        edge_fdim=165,
+        node_fdim=151,
+        num_mt_block=1,
+        num_attn_head=4,
+        depth=6,
+        atom_emb_output="both",
+    )
+    load_grover_weights(model, ckpt_path)
+
+    batch = _make_dummy_batch(num_atoms=5, num_bonds=7, atom_fdim=151, edge_fdim=165)
+    out = model(batch)
+    assert ops.shape(out["atom_from_atom"]) == (5, 800)
+    assert ops.shape(out["bond_from_bond"]) == (7, 800)

k3_node/models/test_jumping_knowledge.py

@@ -0,0 +1,129 @@
+import pytest
+from keras import ops
+
+from k3_node.models import JumpingKnowledge, HeteroJumpingKnowledge
+
+
+def test_jumping_knowledge_cat():
+    num_nodes, channels, num_layers = 100, 17, 5
+    xs = [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+
+    model = JumpingKnowledge('cat')
+    assert str(model) == 'JumpingKnowledge(cat)'
+
+    out = model(xs)
+    assert ops.shape(out) == (num_nodes, channels * num_layers)
+
+
+def test_jumping_knowledge_max():
+    num_nodes, channels, num_layers = 100, 17, 5
+    xs = [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+
+    model = JumpingKnowledge('max')
+    assert str(model) == 'JumpingKnowledge(max)'
+
+    out = model(xs)
+    assert ops.shape(out) == (num_nodes, channels)
+
+
+def test_jumping_knowledge_lstm():
+    num_nodes, channels, num_layers = 100, 17, 5
+    xs = [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+
+    model = JumpingKnowledge('lstm', channels, num_layers)
+    assert str(model) == (
+        f'JumpingKnowledge(lstm, channels={channels}, layers={num_layers})'
+    )
+
+    out = model(xs)
+    assert ops.shape(out) == (num_nodes, channels)
+
+
+def test_jumping_knowledge_invalid_mode():
+    with pytest.raises(AssertionError):
+        JumpingKnowledge('invalid')
+
+
+def test_jumping_knowledge_lstm_missing_args():
+    with pytest.raises(AssertionError):
+        JumpingKnowledge('lstm')
+
+
+def test_hetero_jumping_knowledge_cat():
+    num_nodes, channels, num_layers = 100, 17, 5
+    types = ["author", "paper"]
+    xs_dict = {
+        key: [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+        for key in types
+    }
+
+    model = HeteroJumpingKnowledge(types, mode='cat')
+    model.reset_parameters()
+    assert str(model) == 'HeteroJumpingKnowledge(num_types=2, mode=cat)'
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert ops.shape(out) == (num_nodes, channels * num_layers)
+
+
+def test_hetero_jumping_knowledge_max():
+    num_nodes, channels, num_layers = 100, 17, 5
+    types = ["author", "paper"]
+    xs_dict = {
+        key: [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+        for key in types
+    }
+
+    model = HeteroJumpingKnowledge(types, mode='max')
+    assert str(model) == 'HeteroJumpingKnowledge(num_types=2, mode=max)'
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert ops.shape(out) == (num_nodes, channels)
+
+
+def test_hetero_jumping_knowledge_lstm():
+    num_nodes, channels, num_layers = 50, 8, 3
+    types = ["author", "paper"]
+    xs_dict = {
+        key: [ops.ones((num_nodes, channels)) for _ in range(num_layers)]
+        for key in types
+    }
+
+    model = HeteroJumpingKnowledge(types, mode='lstm', channels=channels,
+                                   num_layers=num_layers)
+    assert str(model) == (
+        f'HeteroJumpingKnowledge(num_types=2, mode=lstm, '
+        f'channels={channels}, layers={num_layers})'
+    )
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert ops.shape(out) == (num_nodes, channels)
+
+
+def test_jumping_knowledge_cat_aggregates_correctly():
+    """Values in cat mode should be the concatenation of all layer reps."""
+    xs = [
+        ops.convert_to_tensor([[1.0, 2.0], [3.0, 4.0]]),
+        ops.convert_to_tensor([[5.0, 6.0], [7.0, 8.0]]),
+    ]
+    model = JumpingKnowledge('cat')
+    out = model(xs)
+    assert ops.shape(out) == (2, 4)
+
+
+def test_jumping_knowledge_max_aggregates_correctly():
+    """Values in max mode should be the element-wise max across layers."""
+    xs = [
+        ops.convert_to_tensor([[1.0, 8.0], [3.0, 4.0]]),
+        ops.convert_to_tensor([[5.0, 6.0], [7.0, 2.0]]),
+    ]
+    import numpy as np
+    model = JumpingKnowledge('max')
+    out = model(xs)
+    assert ops.shape(out) == (2, 2)
+    out_np = ops.convert_to_numpy(out)
+    assert np.allclose(out_np[0], [5.0, 8.0])
+    assert np.allclose(out_np[1], [7.0, 4.0])
+

k3_node/models/test_label_prop.py

@@ -0,0 +1,37 @@
+from keras import ops
+
+from k3_node.models import LabelPropagation
+
+
+def test_label_prop():
+    y = ops.convert_to_tensor([1, 0, 0, 2, 1, 1])
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2, 4, 5], [1, 0, 2, 1, 5, 4]])
+    mask = ops.convert_to_tensor([True, False, True, False, True, False])
+
+    model = LabelPropagation(num_layers=2, alpha=0.5)
+    assert str(model) == 'LabelPropagation(num_layers=2, alpha=0.5)'
+
+    # Without mask:
+    out = model(y, edge_index)
+    assert ops.shape(out) == (6, 3)
+
+    # With mask:
+    out = model(y, edge_index, mask)
+    assert ops.shape(out) == (6, 3)
+
+    # With post_step:
+    out = model(y, edge_index, mask, post_step=lambda y: ops.zeros_like(y))
+    assert float(ops.sum(out)) == 0.0
+
+
+
+def test_label_prop_with_boolean_mask_matches_index_mask():
+    import numpy as np
+
+    y = np.array([0, 1, 2, 0, 1, 2])
+    edge_index = np.array([[0, 1, 2, 3, 4, 5], [3, 4, 5, 0, 1, 2]])
+    model = LabelPropagation(num_layers=3, alpha=0.9)
+    by_bool = model(y, edge_index, mask=np.array([True, True, True, False, False, False]))
+    by_index = model(y, edge_index, mask=np.array([0, 1, 2]))
+    np.testing.assert_allclose(ops.convert_to_numpy(by_bool), ops.convert_to_numpy(by_index), rtol=1e-6)
+    assert (ops.convert_to_numpy(by_bool).argmax(-1) == y).all()

k3_node/models/test_lightgcn.py

@@ -0,0 +1,38 @@
+import pytest
+from keras import ops
+import numpy as np
+
+from k3_node.models import LightGCN
+
+
+@pytest.mark.parametrize('embedding_dim', [16, 32])
+def test_lightgcn(embedding_dim):
+    num_nodes = 50
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5, 6, 7],
+        [1, 2, 3, 4, 5, 6, 7, 0],
+    ])
+    edge_label_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [4, 5, 6, 7],
+    ])
+
+    model = LightGCN(num_nodes, embedding_dim, num_layers=2)
+    assert str(model) == f'LightGCN({num_nodes}, {embedding_dim}, num_layers=2)'
+
+    pred = model(edge_index, edge_label_index)
+    assert ops.shape(pred) == (4,)
+
+    prob = model.predict_link(edge_index, edge_label_index, prob=True)
+    assert ops.shape(prob) == (4,)
+
+    loss = model.recommendation_loss(
+        pos_edge_rank=pred[:2],
+        neg_edge_rank=pred[2:],
+        lambda_reg=1e-4,
+    )
+    assert float(ops.convert_to_numpy(loss)) > 0.0
+
+    recs = model.recommend(edge_index, k=2)
+    assert ops.shape(recs) == (num_nodes, 2)
+

k3_node/models/test_linkx.py

@@ -0,0 +1,31 @@
+import pytest
+from keras import ops
+
+from k3_node.models import LINKX
+
+
+@pytest.mark.parametrize('num_edge_layers', [1, 2])
+def test_linkx(num_edge_layers):
+    x = ops.ones((4, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 3]])
+    edge_weight = ops.convert_to_tensor([0.5, 0.8, 1.0])
+
+    model = LINKX(
+        num_nodes=4,
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=8,
+        num_layers=2,
+        num_edge_layers=num_edge_layers,
+    )
+    assert str(model) == 'LINKX(num_nodes=4, in_channels=16, out_channels=8)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 8)
+
+    out_no_x = model(None, edge_index)
+    assert ops.shape(out_no_x) == (4, 8)
+
+    out_weighted = model(x, edge_index, edge_weight)
+    assert ops.shape(out_weighted) == (4, 8)
+

k3_node/models/test_lpformer.py

@@ -0,0 +1,22 @@
+import numpy as np
+from keras import ops
+
+from k3_node.models.lpformer import LPFormer, get_ppr
+
+
+def test_get_ppr():
+    ppr = get_ppr(np.array([[0, 1, 1, 2, 2, 0], [1, 0, 2, 1, 0, 2]]), num_nodes=3)
+    assert ppr.shape == (3, 3)
+    np.testing.assert_allclose(np.asarray(ppr.sum(axis=1)).ravel(), 1.0, atol=0.05)  # (approximate) distributions
+
+
+def test_lpformer():
+    rng = np.random.default_rng(0)
+    x = rng.random((20, 16)).astype("float32")
+    edge_index = rng.integers(0, 20, (2, 60))
+    edge_index = np.concatenate([edge_index, edge_index[::-1]], axis=1)
+    batch = np.array([[0, 1, 5], [2, 3, 9]])
+    model = LPFormer(in_channels=16, hidden_channels=16, num_gnn_layers=2)
+    ppr = model.calc_sparse_ppr(edge_index, 20)
+    assert tuple(model(batch, x, edge_index, ppr_matrix=ppr).shape) == (3,)
+    assert tuple(model(batch, x, edge_index, ppr_matrix=ppr, training=True).shape) == (3,)

k3_node/models/test_mask_label.py

@@ -0,0 +1,90 @@
+import numpy as np
+from keras import ops
+
+from k3_node.models import MaskLabel
+
+
+def test_mask_label_add():
+    model = MaskLabel(2, 10)
+    assert str(model) == 'MaskLabel()'
+
+    x = ops.ones((4, 10))
+    y = ops.convert_to_tensor([1, 0, 1, 0])
+    mask = ops.convert_to_tensor([False, False, True, True])
+
+    out = model(x, y, mask)
+    assert ops.shape(out) == (4, 10)
+
+    # Non-masked nodes should have same features as input (plus zero embedding)
+    out_np = ops.convert_to_numpy(out)
+    x_np = ops.convert_to_numpy(x)
+    mask_np = ops.convert_to_numpy(mask)
+    # Non-masked embeddings are zero so output matches input exactly
+    assert np.allclose(out_np[~mask_np], x_np[~mask_np])
+
+
+def test_mask_label_concat():
+    model = MaskLabel(2, 10, method='concat')
+
+    x = ops.ones((4, 10))
+    y = ops.convert_to_tensor([1, 0, 1, 0])
+    mask = ops.convert_to_tensor([False, False, True, True])
+
+    out = model(x, y, mask)
+    assert ops.shape(out) == (4, 20)
+
+    # First 10 features should be identical to x
+    out_np = ops.convert_to_numpy(out)
+    x_np = ops.convert_to_numpy(x)
+    assert np.allclose(out_np[:, :10], x_np)
+
+
+def test_mask_label_concat_non_masked_zero():
+    """For concat method, non-masked nodes' label embedding part should be 0."""
+    model = MaskLabel(3, 8, method='concat')
+
+    x = ops.ones((6, 16))
+    y = ops.convert_to_tensor([0, 1, 2, 0, 1, 2])
+    mask = ops.convert_to_tensor([True, True, False, False, False, False])
+
+    out = model(x, y, mask)
+    assert ops.shape(out) == (6, 24)
+
+    out_np = ops.convert_to_numpy(out)
+    mask_np = ops.convert_to_numpy(mask)
+    # Non-masked embedding portion should be all zeros
+    assert np.allclose(out_np[~mask_np, 16:], 0.0)
+
+
+def test_mask_label_invalid_method():
+    import pytest
+    with pytest.raises(ValueError, match="'method' must be either 'add' or 'concat'"):
+        MaskLabel(2, 10, method='multiply')
+
+
+def test_ratio_mask():
+    mask = ops.convert_to_tensor([True, True, True, True, False, False, False, False])
+    out = MaskLabel.ratio_mask(mask, 0.5)
+    out_np = ops.convert_to_numpy(out)
+    mask_np = ops.convert_to_numpy(mask)
+    # True entries in original False positions should stay False
+    assert out_np[~mask_np].sum() == 0
+    # Number of remaining True entries should be <= original count
+    assert out_np[:4].sum() <= 4
+
+
+def test_ratio_mask_zero():
+    """ratio=0 should mask all True entries."""
+    mask = ops.convert_to_tensor([True, True, True, False])
+    out = MaskLabel.ratio_mask(mask, 0.0)
+    out_np = ops.convert_to_numpy(out)
+    assert out_np.sum() == 0
+
+
+def test_ratio_mask_one():
+    """ratio=1.0 should keep all True entries."""
+    mask = ops.convert_to_tensor([True, True, False, False])
+    out = MaskLabel.ratio_mask(mask, 1.0)
+    out_np = ops.convert_to_numpy(out)
+    assert out_np.sum() == 2
+