parameters: (*args, solver_restart: bool = False, **kwargs)
return_statement: return _SR(*args, solver_restart=solver_restart, **kwargs)
docstring: Construct the structure holding the parameters for using the
      Stochastic Reconfiguration/Natural gradient method.
  
      Depending on the arguments, an implementation is chosen. For
      details on all possible kwargs check the specific SR implementations
      in the documentation.
  
      You can also construct one of those structures directly.
  
      Args:
          diag_shift: Diagonal shift added to the S matrix
          method: (cg, gmres) The specific method.
          iterative: Whever to use an iterative method or not.
          jacobian: Differentiation mode to precompute gradients
                    can be "holomorphic", "R2R", "R2C",
                    None (if they shouldn't be precomputed)
          rescale_shift: Whether to rescale the diagonal offsets in SR according
                         to diagonal entries (only with precomputed gradients)
  
      Returns:
          The SR parameter structure.
docstring_summary: Construct the structure holding the parameters for using the
      Stochastic Reconfiguration/Natural gradient method.
docstring_tokens: Construct the structure holding the parameters for using the Stochastic Reconfiguration / Natural gradient method .
function: def build_SR(*args, solver_restart: bool = False, **kwargs):
      """
      Construct the structure holding the parameters for using the
      Stochastic Reconfiguration/Natural gradient method.
  
      Depending on the arguments, an implementation is chosen. For
      details on all possible kwargs check the specific SR implementations
      in the documentation.
  
      You can also construct one of those structures directly.
  
      Args:
          diag_shift: Diagonal shift added to the S matrix
          method: (cg, gmres) The specific method.
          iterative: Whever to use an iterative method or not.
          jacobian: Differentiation mode to precompute gradients
                    can be "holomorphic", "R2R", "R2C",
                    None (if they shouldn't be precomputed)
          rescale_shift: Whether to rescale the diagonal offsets in SR according
                         to diagonal entries (only with precomputed gradients)
  
      Returns:
          The SR parameter structure.
      """
  
      #  try to understand if this is the old API or new
      # API
  
      old_api = False
      # new_api = False
  
      if "matrix" in kwargs:
          # new syntax
          return _SR(*args, **kwargs)
  
      legacy_kwargs = ["iterative", "method"]
      legacy_solver_kwargs = ["tol", "atol", "maxiter", "M", "restart", "solve_method"]
      for key in legacy_kwargs + legacy_solver_kwargs:
          if key in kwargs:
              old_api = True
              break
  
      if len(args) > 0:
          if is_scalar(args[0]):  # it's diag_shift
              old_api = True
          # else:
          #    new_api = True
  
          if len(args) > 1:
              if isinstance(args[1], str):
                  old_api = True
          #     else:
          #        new_api = True
  
      if old_api:
          for (i, arg) in enumerate(args):
              if i == 0:
                  # diag shift
                  kwargs["diag_shift"] = arg
              elif i == 1:
                  kwargs["method"] = arg
              else:
                  raise TypeError(
                      "SR takes at most 2 positional arguments but len(args) where provided"
                  )
  
          args = tuple()
  
          solver = None
          if "iterative" in kwargs:
              kwargs.pop("iterative")
          if "method" in kwargs:
              legacy_solvers = {
                  "cg": jax.scipy.sparse.linalg.cg,
                  "gmres": jax.scipy.sparse.linalg.gmres,
              }
              if kwargs["method"] not in legacy_solvers:
                  raise ValueError(
                      "The old API only supports cg and gmres solvers. "
                      "Migrate to the new API and use any solver from"
                      "jax.scipy.sparse.linalg"
                  )
              solver = legacy_solvers[kwargs["method"]]
              kwargs.pop("method")
          else:
              solver = jax.scipy.sparse.linalg.cg
  
          solver_keys = {}
          has_solver_kw = False
          for key in legacy_solver_kwargs:
              if key in kwargs:
                  solver_keys[key] = kwargs[key]
                  has_solver_kw = True
  
          if has_solver_kw:
              solver = partial(solver, **solver_keys)
  
          kwargs["solver"] = solver
  
      return _SR(*args, solver_restart=solver_restart, **kwargs)
function_tokens: def build_SR ( * args , solver_restart : bool = False , * * kwargs ) : #  try to understand if this is the old API or new # API old_api = False # new_api = False if "matrix" in kwargs : # new syntax return _SR ( * args , * * kwargs ) legacy_kwargs = [ "iterative" , "method" ] legacy_solver_kwargs = [ "tol" , "atol" , "maxiter" , "M" , "restart" , "solve_method" ] for key in legacy_kwargs + legacy_solver_kwargs : if key in kwargs : old_api = True break if len ( args ) > 0 : if is_scalar ( args [ 0 ] ) : # it's diag_shift old_api = True # else: #    new_api = True if len ( args ) > 1 : if isinstance ( args [ 1 ] , str ) : old_api = True #     else: #        new_api = True if old_api : for ( i , arg ) in enumerate ( args ) : if i == 0 : # diag shift kwargs [ "diag_shift" ] = arg elif i == 1 : kwargs [ "method" ] = arg else : raise TypeError ( "SR takes at most 2 positional arguments but len(args) where provided" ) args = tuple ( ) solver = None if "iterative" in kwargs : kwargs . pop ( "iterative" ) if "method" in kwargs : legacy_solvers = { "cg" : jax . scipy . sparse . linalg . cg , "gmres" : jax . scipy . sparse . linalg . gmres , } if kwargs [ "method" ] not in legacy_solvers : raise ValueError ( "The old API only supports cg and gmres solvers. " "Migrate to the new API and use any solver from" "jax.scipy.sparse.linalg" ) solver = legacy_solvers [ kwargs [ "method" ] ] kwargs . pop ( "method" ) else : solver = jax . scipy . sparse . linalg . cg solver_keys = { } has_solver_kw = False for key in legacy_solver_kwargs : if key in kwargs : solver_keys [ key ] = kwargs [ key ] has_solver_kw = True if has_solver_kw : solver = partial ( solver , * * solver_keys ) kwargs [ "solver" ] = solver return _SR ( * args , solver_restart = solver_restart , * * kwargs )