Coverage for src / autoencodix / utils / _model_output.py: 92%

1from typing import Optional 

2from dataclasses import dataclass 

3import torch 

4 

5 

6# internal check done 

7# write tests: done 

8@dataclass 

9class ModelOutput: 

10 """A structured output dataclass for autoencoder models. 

11 

12 This class is used to encapsulate the outputs of autoencoder models in a 

13 consistent format, allowing for flexibility in the type of outputs returned 

14 by different architectures. 

15 

16 Attributes: 

17 reconstruction: The reconstructed input data. 

18 latent_mean: The mean of the latent space distribution, applicable for models like VAEs. 

19 latent_logvar: The log variance of the latent space distribution, applicable for models like VAEs. 

20 additional_info: A dictionary to store any additional information or intermediate outputs. 

21 """ 

22 

23 reconstruction: torch.Tensor 

24 latentspace: torch.Tensor 

25 latent_mean: Optional[torch.Tensor] = None 

26 latent_logvar: Optional[torch.Tensor] = None 

27 additional_info: Optional[dict] = None 

28 

29 def __iter__(self): 

30 yield self