//  [python]
def backward(ctx, grad):
        if ctx.batch_size is not None:
            grad = torch.cat([*grad])

        A, x = ctx.saved_tensors

        # Backprop rule: gradB = A^{-T} grad
        # Check if a workaround for https://github.com/pytorch/pytorch/issues/88890 is needed
        workaround88890 = (
            A.device == torch.device("cpu")
            and (not ctx.upper)
            and ctx.unitriangular
            and (torch.__version__ < (2,))
            and not ctx.transpose
        )
        if not workaround88890:
            gradB = torch.linalg.solve_triangular(
                A, grad, upper=ctx.upper, transpose=not ctx.transpose, unitriangular=ctx.unitriangular
            )
        else:
            # Not sure which workaround it best but for now let's assume we don't want to explicitly transpose A
            # but prefer to let torch.triangular_solve do this inernally
            n = A.shape[0]
            id_csr = torch.sparse_csr_tensor(
                torch.arange(n + 1),
                torch.arange(n),
                torch.ones(n, device=A.device, dtype=A.dtype),
                (n, n),
                device=A.device,
            )
            gradB = torch.linalg.solve_triangular(A + id_csr, grad, upper=ctx.upper, transpose=not ctx.transpose)

        # The gradient with respect to the matrix A seen as a dense matrix would
        # lead to a backprop rule as follows
        # gradA = -(A^{-T} grad)(A^{-1} B) = - gradB @ x.T
        # but we are only interested in the gradient with respect to
        # the (non-zero) values of A. To save memory, instead of computing the full
        # dense matrix gradB @ x.T and then subsampling at the nnz locations in a,
        # we can directly only compute the required values:
        # gradA[i,j] = - dotprod(gradB[i,:], x[j,:])

        # We start by getting the i and j indices:
        A_col_idx = A.col_indices()
        A_crow_idx = A.crow_indices()
        # Uncompress row indices:
        A_row_idx = torch.repeat_interleave(
            torch.arange(A.size()[0], device=A.device), A_crow_idx[1:] - A_crow_idx[:-1]
        )

        if ctx.transpose:
            mgradbselect = -gradB.index_select(0, A_col_idx)  # -gradB[j, :]
            xselect = x.index_select(0, A_row_idx)  # x[i, :]
        else:
            mgradbselect = -gradB.index_select(0, A_row_idx)  # -gradB[i, :]
            xselect = x.index_select(0, A_col_idx)  # x[j, :]

        if ctx.unitriangular is True and torch.any(A_row_idx == A_col_idx):
            raise ValueError("First input should be strictly triangular (i.e. unit diagonals is implicit)")

        # Dot product:
        mgbx = mgradbselect * xselect
        gradA = torch.sum(mgbx, dim=1)

        if ctx.csr is False:
            gradA = torch.sparse_coo_tensor(torch.stack([A_row_idx, A_col_idx]), gradA, A.shape)
        else:
            gradA = torch.sparse_csr_tensor(A_crow_idx, A_col_idx, gradA, A.shape)

        if ctx.batch_size is not None:
            shapes = ctx.A_shape[0] * (ctx.A_shape[-2:],)
            gradA = sparse_block_diag_split(gradA, *shapes)
            if not ctx.csr:
                gradA = torch.stack([*gradA])
            else:
                gradA = stack_csr([*gradA])

            gradB = gradB.view(ctx.B_shape)

        return gradA, gradB, None, None, None