def forward(self, x1, x2, batch_dims=None, **params): x1_ = torch.stack([x1[:, i] for i in self.psi_indices], dim=1) x2_ = torch.stack([x2[:, i] for i in self.psi_indices], dim=1) if batch_dims == (0, 2): print('batch bims here') prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) tone = prod * (self.u1) prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttwo = prod * (self.u2) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 0:1].transpose(-1, -2)) tthree = (diagone + diagtwo) * ((self.rho - 1) * (self.u1)**.5 * (self.u2)**.5) random_effects = tone + ttwo + tthree final = random_effects * self.user_mat final = final + self.first_mat return final
def test_batch_get_indices(self): lhs = torch.randn(2, 5, 1) rhs = torch.randn(2, 1, 5) actual = lhs.matmul(rhs) res = MatmulLazyTensor(lhs, rhs) batch_indices = torch.tensor([0, 1, 0, 1], dtype=torch.long) left_indices = torch.tensor([1, 2, 4, 0], dtype=torch.long) right_indices = torch.tensor([0, 1, 3, 2], dtype=torch.long) self.assertTrue( approx_equal( actual[batch_indices, left_indices, right_indices], res._batch_get_indices(batch_indices, left_indices, right_indices), )) batch_indices = torch.tensor( [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.long) left_indices = torch.tensor( [1, 2, 4, 0, 1, 2, 3, 1, 2, 2, 1, 1, 0, 0, 4, 4, 4, 4], dtype=torch.long) right_indices = torch.tensor( [0, 1, 3, 2, 3, 4, 2, 2, 1, 1, 2, 1, 2, 4, 4, 3, 3, 0], dtype=torch.long) self.assertTrue( approx_equal( actual[batch_indices, left_indices, right_indices], res._batch_get_indices(batch_indices, left_indices, right_indices), ))
def forward(self, x1, x2): if x1.size() == x2.size() and torch.equal(x1, x2): # Use RootLazyTensor when x1 == x2 for efficiency when composing # with other kernels prod = RootLazyTensor(x1 - self.offset) else: prod = MatmulLazyTensor(x1 - self.offset, (x2 - self.offset).transpose(2, 1)) return prod + self.variance.expand(prod.size())
def _get_covariance(self, x1, x2): k_ux1 = self.base_kernel_module(x1, self.inducing_points).evaluate() if torch.equal(x1, x2): covar = RootLazyTensor(k_ux1.matmul(self._inducing_inv_root)) else: k_ux2 = self.base_kernel_module(x2, self.inducing_points).evaluate() covar = MatmulLazyTensor( k_ux1.matmul(self._inducing_inv_root), k_ux2.matmul(self._inducing_inv_root).transpose(-1, -2)) return covar
def test_batch_diag(self): lhs = torch.randn(4, 5, 3) rhs = torch.randn(4, 3, 5) actual = lhs.matmul(rhs) actual_diag = torch.cat([ actual[0].diag().unsqueeze(0), actual[1].diag().unsqueeze(0), actual[2].diag().unsqueeze(0), actual[3].diag().unsqueeze(0), ]) res = MatmulLazyTensor(lhs, rhs) self.assertTrue(approx_equal(actual_diag, res.diag()))
def _get_covariance(self, x1, x2): k_ux1 = delazify(self.base_kernel(x1, self.inducing_points)) if torch.equal(x1, x2): covar = RootLazyTensor(k_ux1.matmul(self._inducing_inv_root)) # Diagonal correction for predictive posterior correction = (self.base_kernel(x1, x2, diag=True) - covar.diag()).clamp(0, math.inf) covar = PsdSumLazyTensor(covar, DiagLazyTensor(correction)) else: k_ux2 = delazify(self.base_kernel(x2, self.inducing_points)) covar = MatmulLazyTensor( k_ux1.matmul(self._inducing_inv_root), k_ux2.matmul(self._inducing_inv_root).transpose(-1, -2)) return covar
def test_matmul(self): lhs = torch.randn(5, 3, requires_grad=True) rhs = torch.randn(3, 4, requires_grad=True) covar = MatmulLazyTensor(lhs, rhs) mat = torch.randn(4, 10) res = covar.matmul(mat) lhs_clone = lhs.clone().detach() rhs_clone = rhs.clone().detach() mat_clone = mat.clone().detach() lhs_clone.requires_grad = True rhs_clone.requires_grad = True mat_clone.requires_grad = True actual = lhs_clone.matmul(rhs_clone).matmul(mat_clone) self.assertTrue(approx_equal(res, actual)) actual.sum().backward() res.sum().backward() self.assertTrue(approx_equal(lhs.grad, lhs_clone.grad)) self.assertTrue(approx_equal(rhs.grad, rhs_clone.grad))
def forward(self, x1, x2, batch_dims=None, **params): action_vector = torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_two]],dim=1) baseline_vector = torch.stack( [torch.Tensor(x1)[:, i] for i in [self.g_indices]], dim=1) fake_vector_one = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) action_vector = torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_two]],dim=1) #combine into new feature vector baseline_vector = torch.stack( [torch.Tensor(x2)[:, i] for i in [self.g_indices]], dim=1) fake_vector_two = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) #x1=[] #print(fake_vector_two) #x1_ = torch.stack([ fake_vector_one[:,i] for i in self.psi_indices],dim=1) #x2_ = torch.stack([fake_vector_two[:,i] for i in self.psi_indices],dim=1) x1_ = fake_vector_one #torch.stack([ fake_vector_one[:,i] for i in self.psi_indices],dim=1) x2_ = fake_vector_two #x1_ = torch.stack([x1[:,i] for i in self.psi_indices],dim=1) #x2_ = torch.stack([x2[:,i] for i in self.psi_indices],dim=1) if batch_dims == (0, 2): print('batch bims here') prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) #.expand(1,100,100) tone = prod * (self.u1) prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttwo = prod * (self.u2) prod = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 2:3].transpose(-1, -2)) tthree = prod * (self.u3) prod = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 3:4].transpose(-1, -2)) tfour = prod * (self.u4) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 0:1].transpose(-1, -2)) cov_12 = (diagone + diagtwo) * ((self.rho_12 - 1) * (self.u1)**.5 * (self.u2)**.5) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 2:3].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 0:1].transpose(-1, -2)) cov_13 = (diagone + diagtwo) * ((self.rho_13 - 1) * (self.u1)**.5 * (self.u3)**.5) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 0:1].transpose(-1, -2)) cov_14 = (diagone + diagtwo) * ((self.rho_14 - 1) * (self.u1)**.5 * (self.u4)**.5) diagone = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 2:3].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 1:2].transpose(-1, -2)) cov_23 = (diagone + diagtwo) * ((self.rho_23 - 1) * (self.u2)**.5 * (self.u3)**.5) diagone = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 1:2].transpose(-1, -2)) cov_24 = (diagone + diagtwo) * ((self.rho_24 - 1) * (self.u2)**.5 * (self.u4)**.5) diagone = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 2:3].transpose(-1, -2)) cov_34 = (diagone + diagtwo) * ((self.rho_34 - 1) * (self.u3)**.5 * (self.u4)**.5) random_effects = tone + ttwo + tthree + tfour + cov_12 + cov_13 + cov_14 + cov_23 + cov_24 + cov_34 #print(random_effects.evaluate()) #print(random_effects) #print(random_effects.size()) #print(self.user_mat.size()) final = random_effects * self.user_mat #print(final.evaluate()) #noise_term = (self.noise**2)*self.noise_mat #print(type(noise_term)) #print(noise_term) #prod = MatmulLazyTensor(x1_, x2_.transpose(-1, -2)) #prod = MatmulLazyTensor(prod,noise_term) #prod = prod*self.user_mat #final = final + noise_term #final = torch.stack((tone,ttwo,tone,ttwo),dim=0) #print('one') #print(random_effects.evaluate()) #print('two') #print(final.evaluate()) #print(MatmulLazyTensor(random_effects,2*torch.eye(100)).evaluate()) #n = self.first_mat #+noise_term final = final + self.first_mat #print(final.evaluate()) return final
def test_diag(self): lhs = torch.randn(5, 3) rhs = torch.randn(3, 5) actual = lhs.matmul(rhs) res = MatmulLazyTensor(lhs, rhs) self.assertTrue(approx_equal(actual.diag(), res.diag()))
def test_transpose(self): lhs = torch.randn(5, 3) rhs = torch.randn(3, 5) actual = lhs.matmul(rhs) res = MatmulLazyTensor(lhs, rhs) self.assertTrue(approx_equal(actual.t(), res.t().evaluate()))
def forward(self, x1, x2, batch_dims=None, **params): action_vector = torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_two]],dim=1) #combine into new feature vector #print('one') baseline_vector =torch.stack([torch.Tensor(x1)[:,i] for i in [self.g_indices]],dim=1) #print('two') #print(baseline_vector.size()) #print(action_vector.size()) fake_vector_one = torch.cat((baseline_vector.squeeze(),action_vector.squeeze()),1) #x1=[] print() #print('three') action_vector = torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_two]],dim=1) #combine into new feature vector baseline_vector =torch.stack([torch.Tensor(x2)[:,i] for i in [self.g_indices]],dim=1) fake_vector_two = torch.cat((baseline_vector.squeeze(),action_vector.squeeze()),1) #x1=[] #fake_vector_one[:,i] x1_ = torch.stack([ fake_vector_one[:,i] for i in self.psi_indices],dim=1) x2_ = torch.stack([fake_vector_two[:,i] for i in self.psi_indices],dim=1) #print(x1_.shape) #print(x2_.shape) #print(x2_[:4,:]) if batch_dims == (0, 2): print('batch bims here') prod = MatmulLazyTensor(x1_[:,0:1], x2_[:,0:1].transpose(-1, -2)) tone = prod * (self.u1) #print('here 1') prod = MatmulLazyTensor(x1_[:,1:2], x2_[:,1:2].transpose(-1, -2)) ttwo = prod * (self.u2) #print('here 2') diagone = MatmulLazyTensor(x1_[:,0:1], x2_[:,1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:,1:2], x2_[:,0:1].transpose(-1, -2)) tthree = (diagone+diagtwo)*((self.rho-1)*(self.u1)**.5*(self.u2)**.5) #print('here 3') random_effects = tone+ttwo+tthree #print(random_effects) final = random_effects*self.user_mat #temp = prod*self.test #temp = temp*self.user_mat #print('here 4') final = final+self.first_mat return final
def forward(self, x1, x2, batch_dims=None, **params): #us = torch.cat([self.u1, self.u2], 0) # us is a vector of size 2 #print(x1[0,:,0:2].size()) # print(x1.size()) #print(us.size()) #x1_ =torch.stack((x1[:,self.psi_dim_one],x1[:,self.psi_dim_two]),dim=1) x1_ = torch.stack([x1[:, i] for i in self.psi_indices], dim=1) #x1_ = torch.stack((x1[:,self.psi_dim_one],x1[:,self.psi_dim_two]),dim=1) #x2_ =torch.stack((x2[:,self.psi_dim_one],x2[:,self.psi_dim_two]),dim=1) x2_ = torch.stack([x2[:, i] for i in self.psi_indices], dim=1) #print(x1_) #print(x2_) #u2_= self.u2 #u1_ =self.u1 #print(self.u1) #print(x1_) #print(x2_) if batch_dims == (0, 2): print('batch bims here') #pass #print(x1_.size()) #x1_ = x1_.view(x1_.size(0), x1_.size(1), -1, 1) #x1_ = x1_.permute(0, 2, 1, 3).contiguous() #x1_ = x1_.view(-1, x1_.size(-2), x1_.size(-1)) #x2_ = x2_.view(x2_.size(0), x2_.size(1), -1, 1) #x2_ = x2_.permute(0, 2, 1, 3).contiguous() #x2_ = x2_.view(-1, x2_.size(-2), x2_.size(-1)) #print(x1_.size()) #print(x2_.size()) #prod = MatmulLazyTensor(x1_, x2_.transpose(1, 0)) prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) #.expand(1,100,100) tone = prod * (self.u1) prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttwo = prod * (self.u2) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 0:1].transpose(-1, -2)) tthree = (diagone + diagtwo) * ((self.rho - 1) * (self.u1)**.5 * (self.u2)**.5) random_effects = tone + ttwo + tthree #print(random_effects.evaluate()) #print(random_effects) #print(random_effects.size()) #print(self.user_mat.size()) final = random_effects * self.user_mat #print(final.evaluate()) #noise_term = (self.noise**2)*self.noise_mat #print(type(noise_term)) #print(noise_term) #prod = MatmulLazyTensor(x1_, x2_.transpose(-1, -2)) #prod = MatmulLazyTensor(prod,noise_term) #prod = prod*self.user_mat #final = final + noise_term #final = torch.stack((tone,ttwo,tone,ttwo),dim=0) #print('one') #print(random_effects.evaluate()) #print('two') #print(final.evaluate()) #print(MatmulLazyTensor(random_effects,2*torch.eye(100)).evaluate()) #n = self.first_mat #+noise_term final = final + self.first_mat #print(final.evaluate()) return final
lambda0 = 0.5 n_cells_1d = 50 forward_cutoff = 400 # Only make 200 observations (Fourier and pointwise). my_problem = ToyFourier2d.build_problem(n_cells_1d, forward_cutoff) updatable_gp = UpdatableGP(kernel, lambda0, sigma0, m0, torch.tensor(my_problem.grid.cells).float(), n_chunks=200) lazy_cov = UpdatableCovLazyTensor(updatable_gp.covariance) # Test getitem. lazy_cov[0:10, 0:10].evaluate() # Test pivoted Cholesky decomposition. from gpytorch.utils.pivoted_cholesky import pivoted_cholesky res = pivoted_cholesky(lazy_cov, max_iter=300, error_tol=0.01) preconditioner = MatmulLazyTensor(res, res.t()) # Now test conjugate gradient inversion. rhs = torch.rand((lazy_cov.n, 1)) ans = linear_cg(lazy_cov.matmul, rhs, tolerance=0.1, max_iter=400, preconditioner=preconditioner.matmul)
def _getitem(self, row_index, col_index, *batch_indices): col_indexer = DiagLazyTensor(torch.ones(self.n))[:, col_index] row_indexer = DiagLazyTensor(torch.ones(self.n))[:, row_index] res = MatmulLazyTensor( MatmulLazyTensor(self, row_indexer).t(), col_indexer) return res
def forward(self, x1, x2, batch_dims=None, **params): #action_vector = torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_one]],dim=1)\ #+torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_two]],dim=1) action_vector = torch.stack( [torch.Tensor(x1)[:, i] for i in [self.action_indices_one]], dim=1) baseline_vector = torch.stack( [torch.Tensor(x1)[:, i] for i in [self.g_indices]], dim=1) #print(action_vector) fake_vector_one = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) #action_vector = torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_one]],dim=1)\ #+torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_two]],dim=1) action_vector = torch.stack( [torch.Tensor(x2)[:, i] for i in [self.action_indices_one]], dim=1) baseline_vector = torch.stack( [torch.Tensor(x2)[:, i] for i in [self.g_indices]], dim=1) fake_vector_two = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) #x1=[] #fake_vector_one[:,i] x1_ = torch.stack([fake_vector_one[:, i] for i in self.psi_indices], dim=1) x2_ = torch.stack([fake_vector_two[:, i] for i in self.psi_indices], dim=1) #print(x1_) if batch_dims == (0, 2): print('batch bims here') prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) tone = prod * (self.u1) #print('here 1') prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttwo = prod * (self.u2) #print('here 2') diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 0:1].transpose(-1, -2)) tthree = (diagone + diagtwo) * ((self.rho - 1) * (self.u1)**.5 * (self.u2)**.5) #print('here 3') random_effects = tone + ttwo + tthree #print(random_effects) final = random_effects * self.user_mat prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) ttimeone = prod * (self.s1) prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttimetwo = prod * (self.s3) time_effects = ttimeone + ttimetwo time_effects = time_effects * self.time_mat final = final + time_effects final = final + self.first_mat return final
def create_lazy_tensor(self): lhs = torch.randn(5, 5, 6, requires_grad=True) rhs = lhs.clone().detach().transpose(-1, -2) covar = MatmulLazyTensor(lhs, rhs) return covar
def create_lazy_tensor(self): lhs = torch.randn(3, 5, 3, requires_grad=True) rhs = torch.randn(3, 3, 6, requires_grad=True) covar = MatmulLazyTensor(lhs, rhs) return covar
def forward(self, x1, x2, batch_dims=None, **params): #x1_ = torch.stack([x1[:,i] for i in self.psi_indices],dim=1) #x2_ = torch.stack([x2[:,i] for i in self.psi_indices],dim=1) action_vector = torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x1)[:,i] for i in [self.action_indices_two]],dim=1) #combine into new feature vector #print('one') print(x1.size()) print('got action vector') #print(action_vector) baseline_vector = torch.stack( [torch.Tensor(x1)[:, i] for i in [self.g_indices]], dim=1) #print('two') #print(baseline_vector.size()) #print(action_vector.size()) fake_vector_one = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) #x1=[] print() #print('three') action_vector = torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_one]],dim=1)\ +torch.stack([torch.Tensor(x2)[:,i] for i in [self.action_indices_two]],dim=1) #combine into new feature vector baseline_vector = torch.stack( [torch.Tensor(x2)[:, i] for i in [self.g_indices]], dim=1) fake_vector_two = torch.cat( (baseline_vector.squeeze(), action_vector.squeeze()), 1) #x1=[] #print(fake_vector_two) #fake_vector_one[:,i] print(self.psi_indices) print(fake_vector_one.size()) print(fake_vector_two.size()) x1_ = fake_vector_one #torch.stack([ fake_vector_one[:,i] for i in self.psi_indices],dim=1) x2_ = fake_vector_two #torch.stack([fake_vector_two[:,i] for i in self.psi_indices],dim=1) print('here') if batch_dims == (0, 2): print('batch bims here') #pass #print(x1_.size()) #x1_ = x1_.view(x1_.size(0), x1_.size(1), -1, 1) #x1_ = x1_.permute(0, 2, 1, 3).contiguous() #x1_ = x1_.view(-1, x1_.size(-2), x1_.size(-1)) #x2_ = x2_.view(x2_.size(0), x2_.size(1), -1, 1) #x2_ = x2_.permute(0, 2, 1, 3).contiguous() #x2_ = x2_.view(-1, x2_.size(-2), x2_.size(-1)) #print(x1_.size()) #print(x2_.size()) #prod = MatmulLazyTensor(x1_, x2_.transpose(1, 0)) prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) #.expand(1,100,100) tone = prod * (self.u1) prod = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 1:2].transpose(-1, -2)) ttwo = prod * (self.u2) prod = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 2:3].transpose(-1, -2)) tthree = prod * (self.u3) prod = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 3:4].transpose(-1, -2)) tfour = prod * (self.u4) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 1:2].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 0:1].transpose(-1, -2)) cov_12 = (diagone + diagtwo) * ((self.rho_12 - 1) * (self.u1)**.5 * (self.u2)**.5) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 2:3].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 0:1].transpose(-1, -2)) cov_13 = (diagone + diagtwo) * ((self.rho_13 - 1) * (self.u1)**.5 * (self.u3)**.5) diagone = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 0:1].transpose(-1, -2)) cov_14 = (diagone + diagtwo) * ((self.rho_14 - 1) * (self.u1)**.5 * (self.u4)**.5) diagone = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 2:3].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 1:2].transpose(-1, -2)) cov_23 = (diagone + diagtwo) * ((self.rho_23 - 1) * (self.u2)**.5 * (self.u3)**.5) diagone = MatmulLazyTensor(x1_[:, 1:2], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 1:2].transpose(-1, -2)) cov_24 = (diagone + diagtwo) * ((self.rho_24 - 1) * (self.u2)**.5 * (self.u4)**.5) diagone = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 3:4].transpose(-1, -2)) diagtwo = MatmulLazyTensor(x1_[:, 3:4], x2_[:, 2:3].transpose(-1, -2)) cov_34 = (diagone + diagtwo) * ((self.rho_34 - 1) * (self.u3)**.5 * (self.u4)**.5) random_effects = tone + ttwo + tthree + tfour + cov_12 + cov_13 + cov_14 + cov_23 + cov_24 + cov_34 #print(random_effects.evaluate()) #print(random_effects) #print(random_effects.size()) #print(self.user_mat.size()) final = random_effects * self.user_mat prod = MatmulLazyTensor(x1_[:, 0:1], x2_[:, 0:1].transpose(-1, -2)) ttimeone = prod * (self.s1**2) prod = MatmulLazyTensor(x1_[:, 2:3], x2_[:, 2:3].transpose(-1, -2)) ttimetwo = prod * (self.s3**2) time_effects = ttimeone + ttimetwo time_effects = time_effects * self.time_mat final = final + time_effects final = final + self.first_mat #print(final.evaluate()) return final