Project: 031_Amortized_Control_of_Continuous_State_Space_Feynma | Function: get_matrix (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the controlled SDE dynamics with locally linear approximation of the drift function, computing matrices A and control α for the linear controlled dynamics

**KEY TERMS:** {
  "exp_A_mat_m": "mean mt",
  "exp_B_mat_m": "control contribution to mean",
  "exp_A_mat_v": "covariance Σt",
  "exp_B_mat_v": "control contribution to covariance",
  "A_mat": "matrix A"
}

**CONTEXT CODE:**     if type(m) == nn.Linear:
        nn.init.orthogonal_(m.weight, 1)        

class LinearSDE(torch.nn.Module):
    def __init__(self, args):
        super(LinearSDE, self).__init__()
        
        self.data = args.dataset
        self.task = args.task
        self.lamda_1 = args.lamda_1
        self.lamda_2 = args.lamda_2
        self.init_sigma = args.init_sigma
        self.ts = args.ts
        self.ld = args.state_dim
        self.nb = args.num_basis
        self.od = args.out_dim
        ### Initial
        
        #### Consturct base matrix for A.
        self.E = nn.Linear(self.ld, self.ld, bias=False)
        self.E.apply(init_orthogonal)
        if GEOTORCH_AVAILABLE:
            geotorch.orthogonal(self.E, "weight")
        else:
            # Fallback: initialize with orthogonal matrix manually
            nn.init.orthogonal_(self.E.weight)
        
        self.D = nn.Parameter(torch.randn(self.nb, self.ld))
        
        ### Init mean and covariance
        self.init_mean = torch.nn.Parameter(torch.randn(self.ld))
        self.init_log_var = torch.nn.Parameter(torch.randn(self.ld))
        self.y_log_var = torch.nn.Parameter(torch.randn(self.ld))
        
        #### Consturct coefficient net for A.
        self.coeff_net = nn.Sequential(nn.Linear(self.ld, self.nb),
                                       nn.Softmax(dim=-1))
        
        ### History encoder
        self.encoder = Transformer_Encoder(args)
        self.decoder = Decoder(args)
        

        self.B = nn.Linear(self.ld, self.ld, bias=False)
        self.B.apply(init_normal)
        self.C = nn.Linear(self.ld, self.ld, bias=False)
        self.C.apply(init_normal)
        self.M = nn.Linear(self.ld, self.ld, bias=False)
        self.M.apply(init_normal)

    def get_matrix(self, alpha, obs_times, sigma=1):
        
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def parallel_compute(self, init, E, Z, obs_times):
        
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def forward(self, obs, obs_times, obs_valid, mask_obs, n_samples=3, epoch=None):
        # For training, we use ELBO in Appendix A.7
        obs_times_ = self.ts * obs_times
        Z, y_observed = self.encoder(obs, obs_times_, obs_mask=mask_obs,  event_mask=obs_valid) # Assume q_phi as dirac delta (do not sampling the latent obs.)
        
        obs_times = obs_times_[..., None]

        E = self.E.weight.data
        init_mean = E.t() @ self.init_mean
        init_var = E.t() @ (self.init_sigma * (elup(self.init_log_var) + 1e-6).diag_embed()) @ E
        init_var = torch.diag(init_var)
        
        init_mean_var = torch.cat([init_mean, init_var], dim=0) 
        
        means, stds, alphas = self.parallel_compute(init_mean_var, self.E, Z, obs_times)
        
        Z = torch.randn(size=(n_samples, *means.size())).to(means.device)
        Y = means + stds * Z

        observed_stamps = obs_valid[..., None]
        
        neg_log_potential = GNLL_(y_observed * observed_stamps, means * observed_stamps, stds**2, eps=1.)
        KLs =  0.5 * (alphas[:, :-1].pow(2) * (obs_times[:, 1:] - obs_times[:, :-1])).mean(0).sum() 
        L_alpha = self.lamda_1 * KLs + self.lamda_2 * neg_log_potential

        mean_out = self.decoder(Y) # Instead of directly decoding y_observed, we decode the samples sampled from controlled latent states.
        var_out = 0.01 * torch.ones_like(mean_out)
        
        return (mean_out, var_out), L_alpha
    

Complete the function implementation replacing "raise NotImplementedError".