Coverage for cosmolayer/cosmolayer/utils.py: 80%

1""" 

2.. module:: cosmolayer.cosmolayer.utils 

3 :synopsis: Utility functions for the COSMO-related computations. 

4 

5.. functionauthor:: Charlles Abreu <craabreu@gmail.com> 

6""" 

7 

8import inspect 

9from collections.abc import Callable 

10from typing import TypeGuard 

11 

12import torch 

13 

14LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] 

15 

16 

17def log_matmul_exp(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: 

18 r"""Compute :math:`\log(\exp(A) \exp(B))` stably in log-space. 

19 

20 Parameters 

21 ---------- 

22 A : torch.Tensor 

23 Tensor of shape (..., M, K). 

24 B : torch.Tensor 

25 Tensor of shape (..., K, N). 

26 

27 Returns 

28 ------- 

29 torch.Tensor 

30 Tensor of shape (..., M, N). 

31 """ 

32 if A.shape[-1] != B.shape[-2]: 

33 raise ValueError("Last dimension of A must match second-to-last dimension of B") 

34 return torch.logsumexp(A.unsqueeze(-1) + B.unsqueeze(-3), dim=-2) 

35 

36 

37def is_loss_function(func: object) -> TypeGuard[LossFn]: 

38 if not callable(func): 

39 return False 

40 

41 try: 

42 sig = inspect.signature(func) 

43 except (TypeError, ValueError): 

44 return False 

45 

46 params = list(sig.parameters.values()) 

47 

48 if len(params) < 2: # noqa: PLR2004 

49 return False 

50 

51 return params[0].name == "input" and params[1].name == "target"