keras/src/ops/linalg.py/Lstsq/compute_output_spec
class Lstsq:    def compute_output_spec(self, a, b):
        if len(a.shape) != 2:
            raise ValueError(
                "Expected a to have rank 2. " f"Received: a.shape={a.shape}"
            )
        if len(b.shape) not in (1, 2):
            raise ValueError(
                "Expected b to have rank 1 or 2. "
                f"Received: b.shape={b.shape}"
            )
        m, n = a.shape
        if b.shape[0] != m:
            raise ValueError(
                "Expected b.shape[0] to be equal to "
                "a.shape[0]. Received: "
                f"a.shape={a.shape}, b.shape={b.shape}"
            )
        if len(b.shape) == 2:
            k = b.shape[1]
            x = KerasTensor((n, k), dtype=a.dtype)
        else:
            x = KerasTensor((n,), dtype=a.dtype)
        return x