submission 824493

import torch

# Helper 1: Calculate Length (Norm)
def calc_length(x: torch.Tensor) -> torch.Tensor:
    """
    Calculates the L2 norm (length) of a batch of vectors.
    We use this to figure out how "long" the column is before we reflect it.
    
    Args:
        x: Tensor of shape [batch_size, vector_length]
    Returns:
        Tensor of shape [batch_size] containing the lengths.
    """
    return torch.norm(x, p=2, dim=1)

# Helper 2: Calculate the Mirror Vector (v) and Scaling Factor (tau)
def calc_mirror_vector(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Builds the Householder reflection vector (the "mirror") for a batch of columns,
    with added protection against rank-deficient (zero) columns.
    """
    # 1. Find the length of the current column
    norm = calc_length(x)
    
    # --- NEW: Safety mask for rank-deficient matrices ---
    # If the norm is incredibly small, we consider the column practically zero.
    is_zero = norm <= 1e-7 
    
    # 2. Get the very first element of each vector in the batch
    x0 = x[:, 0]
    
    # 3. Numerical Stability Trick: 
    sign = torch.sign(x0)
    sign[sign == 0] = 1.0  
    
    # 4. Create the mirror vector 'u'
    u = x.clone()
    u[:, 0] += sign * norm
    
    # 5. The "Compact Format" Requirement:
    u_top = u[:, 0].clone()
    
    # --- NEW: Prevent division by zero ---
    # If the column was zero, fake u_top as 1.0 to prevent NaNs. 
    # (It won't affect the math because we will set tau to 0 later).
    u_top[is_zero] = 1.0 
    
    v = u / u_top.unsqueeze(1)
    
    # 6. Calculate tau (the scaling factor)
    v_sq_sum = torch.sum(v ** 2, dim=1)
    
    # --- NEW: Prevent division by zero for tau ---
    v_sq_sum[is_zero] = 1.0
    
    tau = 2.0 / v_sq_sum
    
    # --- NEW: Turn off the mirror for zero-columns ---
    # If the column was zero, we don't reflect at all.
    tau[is_zero] = 0.0
    
    return v, tau
# Helper 3: Apply the Reflection to the Matrix
def apply_reflection(A: torch.Tensor, v: torch.Tensor, tau: torch.Tensor) -> torch.Tensor:
    """
    Applies the mirror to the remaining parts of the matrix.
    Mathematically: A_new = A - tau * v * (v^T @ A)
    
    Args:
        A: The sub-matrix we are updating. Shape: [batch, rows, cols]
        v: The mirror vector. Shape: [batch, rows]
        tau: The scaling factor. Shape: [batch]
    """
    # Reshape v and tau so PyTorch can do batched matrix multiplication (bmm)
    v_col = v.unsqueeze(2)           # Shape: [batch, rows, 1]
    v_row = v.unsqueeze(1)           # Shape: [batch, 1, rows]
    tau_exp = tau.view(-1, 1, 1)     # Shape: [batch, 1, 1]
    
    # Step A: Project the matrix onto the mirror vector
    # w = v^T @ A
    w = torch.bmm(v_row, A)          # Shape: [batch, 1, cols]
    
    # Step B: Scale it up to create the update matrix
    # update = tau * v * w
    update = tau_exp * torch.bmm(v_col, w)  # Shape: [batch, rows, cols]
    
    # Step C: Subtract the update from the original matrix
    return A - update

# Main Algorithm: Batched QR Factorization
def batched_qr_basic(A_input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    The main loop that walks diagonally down the matrix, zeroing out columns.
    
    Args:
        A_input: The full input batch of square matrices. Shape: [batch, n, n]
        
    Returns:
        H: The packed matrix containing R (top) and v vectors (bottom).
        tau: The 1D tensor of scaling factors.
    """
    A = A_input.clone()
    batch_size, n, _ = A.shape
    
    # Initialize the empty output matrix H with zeros
    H = torch.zeros_like(A)
    
    # We will collect the tau values for each column in this list
    tau_list = []
    
    # Loop through every column EXCEPT the very last one
    # (The last column is a 1x1 block, it's already "upper triangular" by definition)
    for i in range(n - 1):
        
        # 1. Extract the part of the column that is on or below the main diagonal
        x = A[:, i:, i]
        
        # 2. Calculate the mirror vector and tau for this specific column
        v, tau = calc_mirror_vector(x)
        tau_list.append(tau)
        
        # 3. Extract the bottom-right submatrix that needs to be updated
        sub_A = A[:, i:, i:]
        
        # 4. Apply the reflection to update the submatrix
        # This turns the column elements below the diagonal into zeros mathematically
        A[:, i:, i:] = apply_reflection(sub_A, v, tau)
        
        # 5. Pack the 'v' vector into the output matrix 'H'
        # We skip the first element of 'v' (which is exactly 1)
        # We store the rest in the strict lower triangle of H
        H[:, i+1:, i] = v[:, 1:]
        
    # The final column doesn't undergo a reflection, so its tau is just 0
    tau_list.append(torch.zeros(batch_size, dtype=A.dtype, device=A.device))
    
    # Pack the 'R' matrix into 'H'
    # A is now completely upper-triangular. We extract that upper triangle 
    # and add it to H (which currently only holds the lower-triangular v vectors).
    H = H + torch.triu(A)
    
    # Stack our list of 1D tau tensors into a single 2D tensor: [batch, n]
    tau_tensor = torch.stack(tau_list, dim=1)
    
    return H, tau_tensor


# Challenge Entry Point & Testing  

# Type aliases for clarity
input_t = torch.Tensor
output_t = tuple[torch.Tensor, torch.Tensor]

def custom_kernel(data: input_t) -> output_t:
    """
    This is the function the grader will call.
    It replaces the standard torch.geqrf(data).
    """
    return batched_qr_basic(data)