Fix construct_singular_matrix not respecting tags

[patrick-kidger]

No description provided.

[adconner]

Two comments: (1) I think you still dont respect tags. Zeroing the first row matrix = matrix.at[0, :].set(0) is the method for if there are tags, but this does not preserve symmetric matrices, much less psd/nsd matrices (which is the current only usage of tagged singular matrices)

(2) As long as you are thinking about singular matrix construction methods, maybe consider adding some less trivial singular patterns. Eg in #210 for some less trivial tests I change this to (still not attempting to respect tags)

def construct_singular_matrix(
    getkey, solver, tags, num=1, *, size=3, dtype=jnp.float64
):
    # There is currently no attempt to generate matrices respecting tags.
    matrices = construct_matrix(getkey, solver, tags, num, size=size, dtype=dtype)
    if isinstance(solver, (lx.Diagonal, lx.CG, lx.BiCGStab, lx.GMRES)):
        return tuple(matrix.at[0, :].set(0) for matrix in matrices)
    else:
        version = jr.choice(getkey(), np.array([0, 1, 2, 3]))
        if version == 0:
            return tuple(matrix.at[1:, 1:].set(0) for matrix in matrices)
        if version == 1:
            return tuple(matrix.at[0, :].set(0) for matrix in matrices)
        elif version == 2:
            return tuple(matrix.at[:, 0].set(0) for matrix in matrices)
        else:
            return tuple(matrix - matrix.T for matrix in matrices)

If we are serious about respecting tags in singular matrix construction and we also want nontrivial examples we have to think about each tag individually and actually each combination of tags: diagonal and triangular just zero a diagonal element. unit_diagonal impossible when combined with triangular, probably unnecessary when not but we could take any singular matrix with nonzero diagonal and normalize the rows to make the diagonal ones, similar when combined with symmetric or psd or nsd, just do simultaneous row/column normalization by the sqrt of diagonal values. symmetric can zero [1:,1:] to get rank 2 or zero an eigenvalue, psd and nsd probably you have to zero an eigenvalue I don't see a structural option.

[patrick-kidger]

(1) Note that this happens to the matrix before any transformations – for making it respect tags – are applied.

(2) Agreed, it would make sense to add more in here. I think we should be able to simply add more in as random options, or adjust the implementation of _construct_matrix_impl to respect tags+singularity in some way.

This aside, you're seeing this PR a little prematurely! I've not tested this change locally to see if it actually does the right thing, so for now I'm just speculatively trying this out to see what happens in CI.

[adconner]

Ah I see. Only saw the diff should have looked at the rest of the function!

One other restriction I noticed checking this out again is that we have to test matrices with distinct singular values, since otherwise the jax derivative sometimes returns NaN's. In particular tested singular matrices should always have rank exactly 1 less than maximal. So for instance my proposed skew symmetric matrix of odd size doesn't work as there are repeated singular values, causing test failures due to the jax derivative returning NaN