Function test_bitnet_transformer_output_shapes

tests/tests.py:256–260

(
    bitnet_model, input_value, expected_output_shape
)

    ],
)
def test_bitnet_transformer_output_shapes(
    bitnet_model, input_value, expected_output_shape
):
    logits = bitnet_model(input_value)
    assert logits.shape == expected_output_shape


def test_exceptions_on_wrong_dtype():