Code Optimization for For Loops in Task Parameterized Gaussian Mixture Models

Question

I'm currently developing TPGMM from Salinon's work (https://calinon.ch/papers/Calinon-JIST2015.pdf). I'm dealing with large matrix operations using CuPy as shown in the code below. However, I'm having trouble vectorizing the nested for loops in the EM function to make it more efficient. Are there any libraries or algorithms for this sort of vectorization?

Code as attached here

# nbGauss - integer
# nbData - integer
# nbFrames - integer
# nbDim - integer
# nbMinSteps,nbMaxSteps,maxDiffLL,diagRegFact,updateComp = integers
# tpTrajs - nbData x nbDim x nbFrames numpy array
# Mu_tp_init - nbDim x nbFrames x nbGauss numpy array
# Sigma_tp_init - nbDim x nbDim x nbFrames x nbGauss numpy array
# priors_init - nbGauss x 1 numpy array
def fit_em_cupy(nbFrames,nbDim,nbGauss,priors_init,Mu_tp_init,Sigma_tp_init,nbMinSteps,nbMaxSteps,maxDiffLL,diagRegFact,updateComp,tpTrajs):
    """
    Fit the data into the TPGMM model

    nbMinSteps:min number of allowed iterations
    nbMaxSteps:max number of allowed iterations
    maxDiffL:Likelihood increase threshold to stop the algorithm
    diagRegFact:optional regularization
    updateComp:flag to update prior,sigma and mu
    """    
    nbData = tpTrajs.shape[0]
    LL = []
   
    print("Using CuPy")
    # convert to CuPy
    tpTrajs = cp.asarray(tpTrajs)
    priors_fitted = cp.asarray(priors_init)
    Mu_tp_fitted = cp.asarray(Mu_tp_init)
    Sigma_tp_fitted = cp.asarray(Sigma_tp_init)
    for iter in tqdm(range(nbMaxSteps)):

        import time
        now = time.time()
        Lik,priors_fitted,Mu_tp_fitted,Sigma_tp_fitted = EM(nbGauss,nbData,nbDim,nbFrames,tpTrajs,Mu_tp_fitted,Sigma_tp_fitted,priors_fitted,diagRegFact)
        # compute Average Log-likelihood to estimate convergence
        print(f"\nLoop Finished {time.time()-now}")
        Lik_CPU = Lik.get()
        curLL = np.sum(np.log(np.sum(Lik_CPU,0))) / Lik_CPU.shape[1]
        LL.append(curLL)
        if iter>nbMinSteps and ((LL[iter]-LL[iter-1])<maxDiffLL or  iter==nbMaxSteps):
                print(LL[iter])
                print(LL[iter-1])
                print(((LL[iter]-LL[iter-1])/LL[iter-1])*100)
                print(f'EM converged after {iter+1} iterations.')
                return priors_fitted.get(),Mu_tp_fitted.get(),Sigma_tp_fitted.get()
        print(f"Main Finished {time.time()-now}")
    return priors_fitted.get(),Mu_tp_fitted.get(),Sigma_tp_fitted.get()

# nbGauss - integer
# nbData - integer
# nbFrames - integer
# nbDim - integer
# tpTrajs - nbData x nbDim x nbFrames numpy array
# Mu_tp_fitted - nbDim x nbFrames x nbGauss numpy array
# Sigma_tp_fitted - nbDim x nbDim x nbFrames x nbGauss numpy array
# priors_fitted - nbGauss x 1 numpy array
def EM(nbGauss,nbData,nbDim,nbFrames,tpTrajs,Mu_tp_fitted,Sigma_tp_fitted,priors_fitted,diagRegFact):
    
    # E-step
    Lik = cp.ones((nbGauss, nbData))
    GAMMA0 = cp.zeros((nbGauss, nbFrames, nbData))
    for i in range(nbGauss):
        for j in range(nbFrames):
            DataMat = tpTrajs[:,:,j].T - cp.expand_dims(Mu_tp_fitted[:,j,i],axis=-1)
            prob = cp.sum(cp.dot(cp.linalg.inv(Sigma_tp_fitted[:,:,j,i]),DataMat)*DataMat,axis=0)
            GAMMA0[i,j,:] = cp.exp(-0.5*prob) / (cp.sqrt(cp.power(2*cp.pi,nbDim) * cp.abs(cp.linalg.det(Sigma_tp_fitted[:,:,j,i])) + realmin))    
            Lik[i,:] = cp.multiply(Lik[i,:],(GAMMA0[i,j,:]))
        Lik[i,:] = Lik[i,:] * priors_fitted[i]
    GAMMA = Lik / cp.sum(Lik,0)+realmin        
    GAMMA2 = GAMMA / cp.expand_dims(cp.sum(GAMMA,1),axis=-1)

    # M-step
    for i in range(nbGauss):
        # Update Priors
        priors_fitted[i] = cp.sum(GAMMA[i,:]) / nbData

        for j in range(nbFrames):
            DataMat = tpTrajs[:,:,j].T

            # Update Mu
            temp1 = cp.expand_dims(GAMMA2[i,:].T,axis=-1)
            Mu_tp_fitted[:,j,i] = cp.dot(DataMat,temp1)[:,0]

            # update sigma
            DataTmp = DataMat - cp.expand_dims(Mu_tp_fitted[:,j,i],axis=-1)
            Sigma_tp_fitted[:,:,j,i] = cp.dot(DataTmp,cp.dot(cp.diag(GAMMA2[i,:]),DataTmp.T)) + cp.eye(DataTmp.shape[0]) * diagRegFact
    return Lik,priors_fitted,Mu_tp_fitted,Sigma_tp_fitted```