Implement lu_unpack in jax

[barney-s]

Implement lu_unpack in jax

• For Lower and Upper matrices, use jnp.tril, jnp.triu
• Shape both triangle matrices and add 1's to the lower triangle to
  match the pytorch behavior
• For 2D permutation matrix start with an identity matrix and mutate it
  based on pivots
  • start with sequential indices and apply the pivot operations to it
  • finally use the pivoted indices to index the identity matrix to
    generate the final permutation matrix for that pivot.
• For 2D inputs and 1D pivot the above logic would work
• For >=3d inputs, we first reshape the inputs to become 3D and then
  call vmap along the first dim with the 2d logic for each 2d matrix

Fixes: #7507

[qihqi]

this branch looks more correct according to this explanation:

Pivot Matrix Representation

PyTorch uses a compact representation for the permutation matrix. Instead of storing the full matrix, it stores a 1-indexed vector called pivots.

Each element pivots[i] indicates the row that the i-th row was swapped with during the LU decomposition process.
Importantly, it uses 1-indexing, meaning pivots[i] = j signifies that the i-th row was swapped with the (j-1)-th row.
Your Example: [2, 2] -> [[0, 1], [1, 0]]

Let's analyze your example step-by-step:

Input pivots = [2, 2]: This vector tells us:

In the 1st step, the 1st row was swapped with the (2-1) = 1st row (essentially, no swap).
In the 2nd step, the 2nd row was swapped with the (2-1) = 1st row.
Constructing the Permutation Matrix:  We start with an identity matrix and apply the swaps indicated by the pivots vector:

Initial Identity Matrix:

[[1, 0],
 [0, 1]] 
Step 1 (no swap): The matrix remains unchanged.

Step 2 (swap rows 2 and 1):

[[0, 1],
 [1, 0]]

[qihqi]

why transpose here?