Function test_bitnet_transformer_for_various_input_shapes

tests/tests.py:228–231

(bitnet_model, batch_size, seq_len)

    ],
)
def test_bitnet_transformer_for_various_input_shapes(bitnet_model, batch_size, seq_len):
    tokens = torch.randint(0, 20000, (batch_size, seq_len))
    logits = bitnet_model(tokens)
    assert logits.shape == (batch_size, 20000)


def test_rotary_embedding(bitnet_model, random_tensor):