Some simplification patterns can create invalid IR due to type changes

[christopherbate]

What happened?

There are several places in the simplification transforms (e.g. https://github.com/openxla/stablehlo/blob/main/stablehlo/transforms/StablehloAggressiveFolder.cpp#L744) where a rewrite pattern unconditionally replaces a Value with a Value of a potentially different type. Sometimes this is OK, sometimes not.

In this IR, folding the stablehlo.subtract and replacing it directly with a constant of type tensor<1x1xi32> is OK:

func.func @example() -> tensor<?x?xi32> {
    %c = stablehlo.constant dense<1> : tensor<1x1xi32>
    %c1 = stablehlo.constant dense<2> : tensor<1x1xi32>
    %0 = stablehlo.subtract %c, %c1 : (tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<?x?xi32>
    %1 = stablehlo.negate %0 : (tensor<?x?xi32>) -> tensor<?x?xi32>
    return %1 : tensor<?x?xi32>
}

However, if the user of %0 is a return or something like stablehlo.composite, then the result is invalid IR.

It seems like it's hit-or-miss whether patterns are checking and tested against edge cases similar to the above.

Rather than submitting PRs piecemeal each time we encounter a failing pattern, I'm wondering if there's something we can do across the board to address this issue in all the simplification passes.