def plsa_refit_m_step(
    X_rows,
    X_cols,
    X_vals,
    p_w_given_z,
    p_z_given_d,
    p_z_given_wd,
    sample_weight,
    norm_pdz,
):
    """Optimized routine for the M step fitting values of P(z|d) given a fixed set of
    topics (i.e. P(w|z)).

    This routine is optimized to work with sparse matrices and only compute values
    for w, d such that X_{w,d} is non-zero.

    To make this numba compilable the raw arrays defining the COO format sparse
    matrix must be passed separately.

    Parameters
    ----------
    X_rows: array of shape (nnz,)
        For each non-zero entry of X, the row of the entry.

    X_cols: array of shape (nnz,)
        For each non-zero entry of X, the column of the
        entry.

    X_vals: array of shape (nnz,)
        For each non-zero entry of X, the value of entry.

    p_w_given_z: array of shape (n_topics, n_words)
        The fixed topics P(w|z) to fit P(z|d) against.

    p_z_given_d: array of shape (n_docs, n_topics)
        The result array to write new estimates of P(z|d) to.

    p_z_given_wd: array of shape (nnz, n_topics)
        The current estimates for P(z|w,d)

    sample_weight: array of shape (n_docs,)
        Input document weights.

    norm_pdz: array of shape (n_docs,)
        Auxilliary array used for storing row norms; this is passed in to save
        reallocations.

    """

    k = p_z_given_wd.shape[1]
    n = p_z_given_d.shape[0]

    p_z_given_d[:] = 0.0
    norm_pdz[:] = 0.0

    for nz_idx in range(X_vals.shape[0]):
        d = X_rows[nz_idx]
        w = X_cols[nz_idx]
        x = X_vals[nz_idx]

        for z in range(k):
            s = x * p_z_given_wd[nz_idx, z]
            p_z_given_d[d, z] += s
            norm_pdz[d] += s

    for z in range(k):
        for d in range(n):
            if norm_pdz[d] > 0:
                p_z_given_d[d, z] /= norm_pdz[d]

    return p_w_given_z, p_z_given_d