sum_interpolated_lazy_variable.py 文件源码

def diag(self):
        batch_size, n_data, n_interp = self.left_interp_indices.size()

        # Batch compute the non-zero values of the outer products w_left^k w_right^k^T
        left_interp_values = self.left_interp_values.unsqueeze(3)
        right_interp_values = self.right_interp_values.unsqueeze(2)
        interp_values = torch.matmul(left_interp_values, right_interp_values)

        # Batch compute Toeplitz values that will be non-zero for row k
        left_interp_indices = self.left_interp_indices.unsqueeze(3).expand(batch_size, n_data, n_interp, n_interp)
        left_interp_indices = left_interp_indices.contiguous()
        right_interp_indices = self.right_interp_indices.unsqueeze(2).expand(batch_size, n_data, n_interp, n_interp)
        right_interp_indices = right_interp_indices.contiguous()
        batch_interp_indices = Variable(left_interp_indices.data.new(batch_size))
        torch.arange(0, batch_size, out=batch_interp_indices.data)
        batch_interp_indices = batch_interp_indices.view(batch_size, 1, 1, 1)
        batch_interp_indices = batch_interp_indices.expand(batch_size, n_data, n_interp, n_interp).contiguous()
        base_var_vals = self.base_lazy_variable._batch_get_indices(batch_interp_indices.view(-1),
                                                                   left_interp_indices.view(-1),
                                                                   right_interp_indices.view(-1))
        base_var_vals = base_var_vals.view(left_interp_indices.size())

        diag = (interp_values * base_var_vals).sum(3).sum(2).sum(0)
        return diag