@triton.jit
def jagged_2_softmax_kernel(input_ptr, output_ptr, offsets_row_ptr,
    offsets_col_ptr, offsets_overall_ptr, input_stride, output_stride,
    transpose, max_seq_len_row, max_seq_len_col, BLOCK_SIZE: tl.constexpr):
    """
    input shape is [sum_B(Ni * Hi)]
    output shape is [sum_B(Ni * Hi)]
    Padded version = [B, N, H]
    Calculate softmax alone N dim
    Each kernel calulates softmax for 1 sample and 1 head
    offsets_row.size == offsets_col.size == offsets_overall.size
    """
    pid_batch = tl.program_id(0)
    pid_head = tl.program_id(1)
    begin = tl.load(offsets_overall_ptr + pid_batch)
    if transpose:
        N = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load(
            offsets_row_ptr + pid_batch)
        H = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load(
            offsets_col_ptr + pid_batch)
        stride_n = H
        stride_h = H // H
        H = tl.minimum(max_seq_len_col, H)
        N = tl.minimum(max_seq_len_row, N)
    else:
        N = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load(
            offsets_col_ptr + pid_batch)
        H = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load(
            offsets_row_ptr + pid_batch)
        stride_h = N
        stride_n = N // N
        H = tl.minimum(max_seq_len_row, H)
        N = tl.minimum(max_seq_len_col, N)
    if pid_head >= H:
        return
    if H == 0 or N == 0:
        return
    start_ptr = input_ptr + begin * input_stride
    offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = (start_ptr + offsets * input_stride * stride_n + pid_head *
        input_stride * stride_h)
    row = tl.load(input_ptrs, mask=offsets < N, other=-float('inf'))
    row_mins_max = row - tl.max(row, axis=0)
    numerator = tl.exp(row_mins_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator
    output_start_ptr = output_ptr + begin * output_stride
    output_ptrs = (output_start_ptr + offsets * output_stride * stride_n + 
        pid_head * output_stride * stride_h)
    tl.store(output_ptrs, softmax_output, mask=offsets < N)
