class_docstr: Class for NOTEARS MLP (Multi-layer Perceptron) model.
  The model weights consist of dag_layer and loc_lin_layer weights respectively.
  dag_layer weight is the weight of the first fully connected layer which determines the causal structure.
  loc_lin_layer weights are the weight of hidden layers after the first fully connected layer
class_docstr_tok: ["Class","for","NOTEARS","MLP","(","Multi-layer","Perceptron",")","model",".","The","model","weights","consist","of","dag_layer","and","loc_lin_layer","weights","respectively",".","dag_layer","weight","is","the","weight","of","the","first","fully","connected","layer","which","determines","the","causal","structure",".","loc_lin_layer","weights","are","the","weight","of","hidden","layers","after","the","first","fully","connected","layer"]
docstr: Perform one step of dual ascent in augmented Lagrangian.
  
  Args:
      X: input tensor data.
      rho: max number of dual ascent steps during optimisation.
      alpha: exit if h(w) < h_tol (as opposed to strict definition of 0).
      h: DAGness of the adjacency matrix
      rho_max: to be updated
  
  Returns:
      rho, alpha and h
docstr_tok: ["Perform","one","step","of","dual","ascent","in","augmented","Lagrangian",".","Args",":","X",":","input","tensor","data",".","rho",":","max","number","of","dual","ascent","steps","during","optimisation",".","alpha",":","exit","if","h","(","w",")","<","h_tol","(","as","opposed","to","strict","definition","of","0",")",".","h",":","DAGness","of","the","adjacency","matrix","rho_max",":","to","be","updated","Returns",":","rho",",","alpha","and","h"]
code: def _dual_ascent_step(
      self,
      X: torch.Tensor,
      rho: float,
      alpha: float,
      h: float,
      rho_max: float,
  ) -> Tuple[float, float, float]:
      """
      Perform one step of dual ascent in augmented Lagrangian.
  
      Args:
          X: input tensor data.
          rho: max number of dual ascent steps during optimisation.
          alpha: exit if h(w) < h_tol (as opposed to strict definition of 0).
          h: DAGness of the adjacency matrix
          rho_max: to be updated
  
      Returns:
          rho, alpha and h
      """
  
      def _get_flat_grad(params: List[torch.Tensor]) -> np.ndarray:
          """
          Get flatten gradient vector from the parameters of the model
  
          Args:
              params: parameters of the model
  
          Returns:
              flatten gradient vector in numpy form
          """
          views = [
              p.data.new(p.data.numel()).zero_()
              if p.grad is None
              else p.grad.data.to_dense().view(-1)
              if p.grad.data.is_sparse
              else p.grad.data.view(-1)
              for p in params
          ]
          return torch.cat(views, 0).cpu().detach().numpy()
  
      def _get_flat_bounds(
          params: List[torch.Tensor],
      ) -> List[Tuple[Union[None, float]]]:
          """
          Get bound constraint for each parameter in flatten vector form from the parameters of the model
  
          Args:
              params: parameters of the model
  
          Returns:
              flatten vector of bound constraints for each parameter in numpy form
          """
          bounds = []
  
          for p in params:
              try:
                  b = p.bounds
              except AttributeError:
                  b = [(None, None)] * p.numel()
              bounds += b
  
          return bounds
  
      def _get_flat_params(params: List[torch.Tensor]) -> np.ndarray:
          """
          Get parameters in flatten vector from the parameters of the model
  
          Args:
              params: parameters of the model
  
          Returns:
              flatten parameters vector in numpy form
          """
          views = [
              p.data.to_dense().view(-1)
              if p.data.is_sparse
              else p.data.view(-1)
              for p in params
          ]
          return torch.cat(views, 0).cpu().detach().numpy()
  
      def _update_params_from_flat(
          params: List[torch.Tensor], flat_params: np.ndarray
      ):
          """
          Update parameters of the model from the parameters in the form of flatten vector
  
          Args:
              params: parameters of the model
              flat_params: parameters in the form of flatten vector
          """
          offset = 0
          flat_params_torch = (
              torch.from_numpy(flat_params)
              .to(torch.get_default_dtype())
              .to(self.device)
          )
          for p in params:
              n_params = p.numel()
              # view_as to avoid deprecated pointwise semantics
              p.data = flat_params_torch[
                  offset : offset + n_params
              ].view_as(p.data)
              offset += n_params
  
      def _func(flat_params: np.ndarray) -> Tuple[float, np.ndarray]:
          """
          Objective function that the NOTEARS algorithm tries to minimise.
  
          Args:
              flat_params: parameters to be optimised to minimise the objective function
  
          Returns:
              Loss and gradient
          """
          _update_params_from_flat(params, flat_params)
          optimizer.zero_grad()
  
          n_features = X.shape[1]
          X_hat = self(X)
          h_val = self._h_func()
          loss = 0.0
  
          # sum the losses across all dist types
          for dist_type in self.dist_types:
              loss = loss + dist_type.loss(X, X_hat)
  
          lagrange_penalty = 0.5 * rho * h_val * h_val + alpha * h_val
          # NOTE: both the l2 and l1 regularization are NOT applied to the bias parameters
          l2_reg = 0.5 * self.ridge_beta * self._l2_reg(n_features)
          l1_reg = self.lasso_beta * self._l1_reg(n_features)
  
          primal_obj = loss + lagrange_penalty + l2_reg + l1_reg
          primal_obj.backward()
          loss = primal_obj.item()
  
          flat_grad = _get_flat_grad(params)
          return loss, flat_grad.astype("float64")
  
      optimizer = torch.optim.Optimizer(self.parameters(), {})
      params = optimizer.param_groups[0]["params"]
      flat_params = _get_flat_params(params)
      bounds = _get_flat_bounds(params)
      h_new = np.inf
  
      while (rho < rho_max) and (h_new > 0.25 * h or h_new == np.inf):
          # Magic
          sol = sopt.minimize(
              _func,
              flat_params,
              method="L-BFGS-B",
              jac=True,
              bounds=bounds,
          )
          _update_params_from_flat(params, sol.x)
          h_new = self._h_func().item()
  
          if h_new > 0.25 * h:
              rho *= 10
  
      alpha += rho * h_new
      return rho, alpha, h_new
  
imports: import logging
  from typing import Iterable, List, Tuple, Union
  import numpy
  import scipy.optimize
  import torch
  from sklearn.base import BaseEstimator
  from torch import nn
  from causalnex.structure.pytorch.dist_type._base import DistTypeBase
  from causalnex.structure.pytorch.nonlinear import LocallyConnected
label_desc: Use image node_id 10 for calling the NotearsMLP obj's underlying member method code with example usage: obj._dual_ascent_step(X, rho, alpha, h, rho_max) and returns: rho,  alpha,  h_new, torch, bounds, torch, loss,  flat_grad