import torch
import torch.nn as nn

BATCH_SIZE = 256
DIM = 128 
SHAPE = (BATCH_SIZE, DIM, DIM)

# 使用更高的数据类型以保证精度
DTYPE = torch.float64

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return torch.linalg.cholesky(input_tensor)

def get_inputs():
    """
    生成用于测试的输入张量。
    输入必须是对称正定(Symmetric Positive-Definite)矩阵。
    一个简单的方法是 A = B @ B.T + eps * I
    """
    # Create a random matrix B
    B = torch.randn(BATCH_SIZE, DIM, DIM, dtype=DTYPE)
    # Construct A = B @ B.T to make it symmetric and positive semi-definite
    A = torch.bmm(B, B.transpose(1, 2))
    # Add a small value to the diagonal to ensure it's strictly positive-definite
    # and improve numerical stability.
    A.diagonal(offset=0, dim1=-1, dim2=-2).add_(1e-10)
    
    return [A.contiguous()]

def get_init_inputs():
    return []