keras/src/backend/tensorflow/linalg.py/solve_triangular
def solve_triangular(a, b, lower=False):
    if b.shape.ndims == a.shape.ndims - 1:
        b = tf.expand_dims(b, axis=-1)
        return tf.squeeze(
            tf.linalg.triangular_solve(a, b, lower=lower), axis=-1
        )
    return tf.linalg.triangular_solve(a, b, lower=lower)