import unittest
from resource_estimates import dense_transformer_parameters, training_flops, weight_memory_bytes


class ResourceEstimateTests(unittest.TestCase):
    def test_training_flops_uses_6nd_estimate(self):
        self.assertEqual(training_flops(3_000_000_000, 500_000_000_000), 9_000_000_000_000_000_000_000)

    def test_int4_weight_memory_is_quarter_of_fp16(self):
        parameters = 8_000_000_000
        self.assertEqual(weight_memory_bytes(parameters, 4) * 4, weight_memory_bytes(parameters, 16))

    def test_parameter_estimate_rejects_invalid_shapes(self):
        with self.assertRaises(ValueError):
            dense_transformer_parameters(0, 512, 2048, 32000)


if __name__ == "__main__":
    unittest.main()