forked from ccf-ai-infra/TileOPs-Metax
19 lines
624 B
Python
19 lines
624 B
Python
import torch
|
|
|
|
|
|
def get_tolerances(dtype: torch.dtype) -> dict:
|
|
"""Return atol/rtol dict for GLA correctness tests."""
|
|
if dtype == torch.float32:
|
|
return {"atol": 1e-2, "rtol": 1e-2}
|
|
elif dtype == torch.float16:
|
|
return {"atol": 5e-2, "rtol": 5e-2}
|
|
else: # bfloat16
|
|
return {"atol": 1e-1, "rtol": 1e-1}
|
|
|
|
|
|
def cosine_sim(a: torch.Tensor, b: torch.Tensor) -> float:
|
|
"""Compute cosine similarity between two tensors (flattened)."""
|
|
a_flat = a.float().flatten()
|
|
b_flat = b.float().flatten()
|
|
return (torch.dot(a_flat, b_flat) / (a_flat.norm() * b_flat.norm() + 1e-12)).item()
|