Write a custom CUDA kernel to optimize `torch.cholesky` (torch.linalg.cholesky).

The operation performs Cholesky decomposition on a batch of symmetric positive-definite matrices. It decomposes a matrix `A` into `L * L.T`, where `L` is a lower triangular matrix.

**Problem Analysis:**
Cholesky decomposition is a complex algorithm with strong data dependencies, unlike simple element-wise or GEMM operations.
1.  **Sequential Dependency**: To compute column `j` of the output matrix `L`, all preceding columns `k < j` must already be computed. This introduces an inherent sequential nature to the algorithm.
2.  **Intra-Column Dependency**: Within the computation of column `j`, the diagonal element `L[j,j]` must be calculated first, as all other elements `L[i,j]` in that column depend on it.
3.  **Library Optimization**: PyTorch's implementation calls highly specialized and optimized routines from libraries like MAGMA or cuSOLVER. A custom kernel must be intelligently designed to manage the dependencies to be competitive.

**Optimization Strategy: Batched Kernel with Hybrid Parallel-Sequential Logic**

The strategy is to create a single CUDA kernel that handles a batch of matrices. For each matrix, it uses a sequential outer loop over the columns, but executes the computations within each step in parallel.

1.  **Block-per-Matrix Parallelism**: The kernel is launched with a grid of blocks, where each thread block is responsible for decomposing a single matrix from the input batch.

2.  **Outer Sequential Loop**: Inside the kernel, a `for` loop iterates from column `j = 0` to `n-1`. This loop respects the primary data dependency of the algorithm.

3.  **Inner Parallel Computation Stages**: Within each iteration of the loop, the work is parallelized:
    *   **Stage 1: Parallel Reduction for Diagonal Element**: To compute `L[j,j]`, a dot product is required. The threads in the block perform this as a parallel reduction using **shared memory** for fast data exchange and summation. The first thread then computes the final `sqrt` and writes the result.
    *   **Synchronization**: A `__syncthreads()` barrier ensures the diagonal element is visible to all threads.
    *   **Stage 2: Parallel Computation for Off-Diagonal Elements**: The threads in the block then work in parallel to compute all other elements `L[i,j]` in the current column. Each thread is assigned one or more rows `i > j` and computes the necessary dot products and final value independently.

This approach correctly handles the algorithm's dependencies while maximizing parallelism at each stage, encapsulating the entire complex decomposition within a single, self-contained kernel.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
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 []