from __future__ import annotations

import unittest

from airfrans_frontier.runtime import remove_pythonpath_entries

remove_pythonpath_entries()

import torch

from airfrans_frontier.models import PointwiseMLP


class PointwiseMLPTests(unittest.TestCase):
    def test_forward_pass_returns_batch_by_target_dim(self) -> None:
        model = PointwiseMLP(input_dim=5, output_dim=4, hidden_width=16, depth=2, activation="gelu")
        batch = torch.randn(7, 5)

        output = model(batch)

        self.assertEqual(tuple(output.shape), (7, 4))


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