import torch
import torch.nn as nn

from trace_commentor.parser import *


def test_func_no_arg():

    @analyse
    def target():

        x = torch.ones(4, 5)
        for i in range(3):
            x = x[..., None, :]

        a = torch.randn(309, 110, 3)[:100]
        f = nn.Linear(3, 128)
        b = f(a.reshape(-1, 3)).reshape(309, 110, 128)
        c = torch.concat((a, b), dim=-1)

        return c.flatten()
    
    print()
    target()


def test_for_loop():

    @analyse
    def target():
        a = 1
        for i in range(3):
            a += 1
        print(a)
    
    print()
    target()