Batched TriangularSolve of singular matrix returns incorrect results

[jakevdp]

Here's a short repro in JAX, that more or less passes the inputs directly to TriangularSolveOp:

import jax
import jax.numpy as jnp

def solve(x, y):
  return jax.lax.linalg.triangular_solve(x, y, left_side=True)

x = jnp.array([[1., 1.], [0., 0.]])
y = jnp.array([[1], [1.]])

print(solve(x, y))
# [[-inf]
#  [ inf]]
print(solve(x[None], y[None])[0])
# [[0.]
#  [1.]]

I would expect the second output to be identical to the first. This appears to be the root cause of the issue reported in jax-ml/jax#15429

[Qiustander]

Is this issue coming from the batching rule of triangular_solve? Had it solved?

[hawkinsp]

Fixed by #15026