perforatedai.globals_perforatedai

   1# Copyright (c) 2025 Perforated AI
   2"""PAI configuration file.
   3
   4This module provides configuration classes and utilities for Perforated AI (PAI),
   5including device settings, dendrite management, module conversion options,
   6and training parameters.
   7"""
   8
   9import math
  10import sys
  11
  12import torch
  13import torch.nn as nn
  14
  15
  16def _validate_module_id(module_id):
  17    """Validate that a module ID string uses dot notation.
  18
  19    Module IDs must start with '.' and must not contain '[' or ']'.
  20    For example, 'model.layers[1].module' should be written as '.layers.1.module'.
  21    """
  22    if not isinstance(module_id, str) or not module_id.startswith("."):
  23        raise ValueError(
  24            f"Module ID '{module_id}' must start with '.' - model.module should be '.module'"
  25        )
  26    if "[" in module_id or "]" in module_id:
  27        raise ValueError(
  28            f"Module ID '{module_id}' must not contain '[' or ']'. "
  29            "Use dot notation instead, e.g. '.layers.1.module'"
  30        )
  31
  32
  33def add_pai_config_var_functions(obj, var_name, initial_value, list_type=False):
  34    """Dynamically add a property with getter and setter to an object.
  35
  36    This function adds a private variable along with getter and setter methods
  37    to a given object instance. Used for integrating initial and Perforated
  38    Backpropagation variables into the PAIConfig class.
  39
  40    Parameters
  41    ----------
  42    obj : object
  43        The object to which the property will be added.
  44    var_name : str
  45        Name of the variable/property to create.
  46    initial_value : any
  47        Initial value for the property.
  48
  49    Returns
  50    -------
  51    None
  52
  53    Notes
  54    -----
  55    Creates three attributes on obj:
  56        - _{var_name}: private storage
  57        - get_{var_name}: getter method
  58        - set_{var_name}: setter method
  59    """
  60    private_name = f"_{var_name}"
  61
  62    # Add the private variable to the instance
  63    setattr(obj, private_name, initial_value)
  64
  65    # Define getter and setter and appender
  66
  67    def getter_val(self):
  68        """Get the current value of the property.
  69
  70        If the property a individual value but is set to be a list,
  71        return the element corresponding to the
  72        current number of dendrites added. Otherwise, return the value directly.
  73
  74        Returns
  75        -------
  76        Any
  77            Current value of the property.
  78
  79        Notes:
  80        -----
  81        Many variables have optimal settings that must change as dendrites are added
  82        this enables those values to be dynamically set very easily.
  83
  84        Parameters
  85        ----------
  86        None
  87
  88        """
  89        global pai_tracker
  90        if type(getattr(self, private_name)) is list:
  91            return getattr(self, private_name)[
  92                min(
  93                    len(getattr(self, private_name)) - 1,
  94                    pai_tracker.member_vars["num_dendrites_added"],
  95                )
  96            ]
  97        return getattr(self, private_name)
  98
  99    def getter_list(self):
 100        """Get the underlying list value for this dynamic config property.
 101
 102        Returns
 103        -------
 104        list
 105            Raw list stored for this property.
 106
 107        Parameters
 108        ----------
 109        None
 110
 111        """
 112        return getattr(self, private_name)
 113
 114    def setter(self, value):
 115        """Set the value of the property.
 116
 117        Parameters
 118        ----------
 119        value : Any
 120            New value to assign to the dynamic configuration field.
 121
 122        Returns
 123        -------
 124        None
 125            Updates the property and may trigger configuration auto-save.
 126        """
 127        if (
 128            self.__dict__.get("_module_name") is not None
 129            or self.__dict__.get("_module_type") is not None
 130        ):
 131            raise RuntimeError(
 132                "Setting custom module config values should only be done "
 133                "from JSON config files or the GUI"
 134            )
 135        if var_name in (
 136            "module_ids_to_track",
 137            "module_ids_to_perforate",
 138            "parameter_ids_to_track",
 139        ):
 140            for module_id in value:
 141                _validate_module_id(module_id)
 142        setattr(self, private_name, value)
 143        if not self.__dict__.get("_loading_config_values", False):
 144            self.__dict__.setdefault("_manually_set_keys", set()).add(var_name)
 145            if var_name == "config_file":
 146                self.sync_config_sources()
 147
 148    def appender(self, value):
 149        """Append a value to the property if it is a list.
 150
 151        Parameters
 152        ----------
 153        value : Any
 154            Value or values to append to the underlying list property.
 155
 156        Returns
 157        -------
 158        None
 159            Appends items in-place and prints the resulting list value.
 160        """
 161        if isinstance(getattr(self, private_name), list):
 162            if var_name in (
 163                "module_ids_to_track",
 164                "module_ids_to_perforate",
 165                "parameter_ids_to_track",
 166            ):
 167                for module_id in value:
 168                    _validate_module_id(module_id)
 169            setattr(self, private_name, getattr(self, private_name) + value)
 170            print(
 171                'New list value of "{}": {}'.format(
 172                    private_name, getattr(self, private_name)
 173                )
 174            )
 175        else:
 176            raise TypeError(f"Cannot append to non-list attribute '{var_name}'")
 177
 178    # Attach methods to the instance
 179    if list_type:
 180        setattr(obj, f"get_{var_name}", getter_list.__get__(obj))
 181    else:
 182        setattr(obj, f"get_{var_name}", getter_val.__get__(obj))
 183    setattr(obj, f"set_{var_name}", setter.__get__(obj))
 184    setattr(obj, f"append_{var_name}", appender.__get__(obj))
 185
 186
 187# ---------------------------------------------------------------------------
 188# JSON serialization helpers  (used by PAIConfig.save_config / load_config)
 189# ---------------------------------------------------------------------------
 190
 191
 192def _resolve_dotted_name(dotted_name):
 193    """Import and return an object identified by a dotted module path.
 194
 195    E.g. 'torch.nn.modules.conv.Conv2d'  → nn.Conv2d class
 196         'torch.sigmoid'                 → torch.sigmoid function
 197
 198    Returns None if the name cannot be resolved.
 199    """
 200    import importlib
 201
 202    # Common shorthand names accepted in user-facing config files.
 203    if dotted_name == "sigmoid":
 204        return torch.sigmoid
 205    if dotted_name == "relu":
 206        return torch.relu
 207    if dotted_name == "tanh":
 208        return torch.tanh
 209
 210    parts = dotted_name.rsplit(".", 1)
 211    if len(parts) == 2:
 212        try:
 213            mod = importlib.import_module(parts[0])
 214            return getattr(mod, parts[1], None)
 215        except Exception:
 216            pass
 217    # Fall back: try the whole string as a single attribute of builtins
 218    import builtins
 219
 220    return getattr(builtins, dotted_name, None)
 221
 222
 223_TORCH_ACTIVATION_SHORTHAND = ("sigmoid", "relu", "tanh")
 224
 225
 226def callable_config_repr(val):
 227    """Canonical string form of a callable config value (e.g. pai_forward_function).
 228
 229    Shared by config serialisation and the interactive configuration display so
 230    the two never disagree on how a callable is named.
 231    """
 232    name = getattr(val, "__name__", None) or getattr(val, "__qualname__", None)
 233    mod = getattr(val, "__module__", None)
 234    if name and mod:
 235        if name in _TORCH_ACTIVATION_SHORTHAND:
 236            return f"torch.{name}"
 237        return f"{mod}.{name}"
 238    if name:
 239        return str(name)
 240    return repr(val)
 241
 242
 243def _serialize_pai_value(val):
 244    """Recursively convert a PAIConfig value to a JSON-serialisable form."""
 245    if isinstance(val, bool):
 246        return val
 247    if isinstance(val, (int, float, str)):
 248        return val
 249    if isinstance(val, torch.device):
 250        return str(val)
 251    if isinstance(val, torch.dtype):
 252        return str(val)
 253    if isinstance(val, list):
 254        return [_serialize_pai_value(v) for v in val]
 255    if isinstance(val, type):
 256        mod = getattr(val, "__module__", "") or ""
 257        return f"{mod}.{val.__name__}" if mod else val.__name__
 258    if callable(val):
 259        return callable_config_repr(val)
 260    return str(val)
 261
 262
 263def _deserialize_pai_value(json_val, type_hint):
 264    """Convert a JSON value to its Python type using an explicit type hint."""
 265    if json_val is None:
 266        return None
 267    if type_hint is torch.device:
 268        return torch.device(str(json_val))
 269    if type_hint is torch.dtype:
 270        v = getattr(torch, str(json_val).split(".", 1)[-1], None)
 271        return v if isinstance(v, torch.dtype) else json_val
 272    if type_hint is callable:
 273        v = _resolve_dotted_name(json_val) if isinstance(json_val, str) else None
 274        return v if (v and callable(v)) else json_val
 275    if type_hint == [type]:  # list whose elements are class objects
 276        return [
 277            (_resolve_dotted_name(v) if isinstance(v, str) else v)
 278            for v in (json_val or [])
 279        ]
 280    return json_val  # bool, int, float, str, list — JSON value is already correct
 281
 282
 283class PAIConfig:
 284    """Configuration class for PAI settings.
 285
 286    This class manages all configuration parameters for the Perforated AI system,
 287    including device settings, dendrite behavior, module conversion rules,
 288    training parameters, and debugging options.
 289
 290    Attributes
 291    ----------
 292    use_cuda : bool
 293        Whether CUDA is available and should be used.
 294    device : torch.device
 295        The device to use for computation (CPU, CUDA, etc.).
 296    save_name : str
 297        Name used for saving models (should not be set manually).
 298    debugging_output_dimensions : int
 299        Debug level for input dimension checking.
 300    confirm_correct_sizes : bool
 301        Whether to verify tensor sizes during execution.
 302    unwrapped_modules_confirmed : bool
 303        Confirmation flag for using unwrapped modules.
 304    weight_decay_accepted : bool
 305        Confirmation flag for accepting weight decay.
 306    checked_skipped_modules : bool
 307        Whether skipped modules have been verified.
 308    verbose : bool
 309        Enable verbose logging output.
 310    extra_verbose : bool
 311        Enable extra verbose logging output.
 312    silent : bool
 313        Suppress all PAI print statements.
 314    save_old_graph_scores : bool
 315        Whether to save historical graph scores.
 316    testing_dendrite_capacity : bool
 317        Enable dendrite capacity testing mode.
 318    using_safe_tensors : bool
 319        Use safe tensors file format for saving.
 320    global_candidates : int
 321        Number of global candidate dendrites.
 322    drawing_pai : bool
 323        Enable PAI visualization graphs.
 324    test_saves : bool
 325        Save intermediary test models.
 326    pai_saves : bool
 327        Save PAI-specific format models.
 328    output_dimensions : list
 329        Format specification for input tensor dimensions.
 330    improvement_threshold : float
 331        Relative improvement threshold for validation scores.
 332    improvement_threshold_raw : float
 333        Absolute improvement threshold for validation scores.
 334    candidate_weight_initialization_multiplier : float
 335        Multiplier for random dendrite weight initialization.
 336    DOING_SWITCH_EVERY_TIME : int
 337        Constant for switch mode: add dendrites every epoch.
 338    DOING_HISTORY : int
 339        Constant for switch mode: add dendrites based on validation history.
 340    n_epochs_to_switch : int
 341        Number of epochs without improvement before switching.
 342    history_lookback : int
 343        Number of epochs to average for validation history.
 344    initial_history_after_switches : int
 345        Epochs to wait after adding dendrites before beggining checks.
 346    DOING_FIXED_SWITCH : int
 347        Constant for switch mode: add dendrites at fixed intervals.
 348    fixed_switch_num : int
 349        Number of epochs between fixed switches.
 350    first_fixed_switch_num : int
 351        Number of epochs before first switch (for pretraining).
 352    DOING_NO_SWITCH : int
 353        Constant for switch mode: never add dendrites.
 354    switch_mode : int
 355        Current switch mode setting.
 356    reset_best_score_on_switch : bool
 357        Whether to reset best score when adding dendrites.
 358    learn_dendrites_live : bool
 359        Enable live dendrite learning (advanced feature).
 360    no_extra_n_modes : bool
 361        Disable extra neuron modes (advanced feature).
 362    d_type : torch.dtype
 363        Data type for dendrite weights.
 364    retain_all_dendrites : bool
 365        Keep dendrites even if they don't improve performance.
 366    find_best_lr : bool
 367        Automatically sweep learning rates when adding dendrites.
 368    dont_give_up_unless_learning_rate_lowered : bool
 369        Ensure search lowers learning rate at least once.
 370    max_dendrite_tries : int
 371        Maximum attempts to add dendrites with random initializations.
 372    max_dendrites : int
 373        Maximum total number of dendrites to add.
 374    PARAM_VALS_BY_TOTAL_EPOCH : int
 375        Constant: scheduler params tracked by total epochs.
 376    PARAM_VALS_BY_UPDATE_EPOCH : int
 377        Constant: scheduler params reset at each switch.
 378    PARAM_VALS_BY_NEURON_EPOCH_START : int
 379        Constant: scheduler params reset for neuron starts only.
 380    param_vals_setting : int
 381        Current parameter tracking mode.
 382    pai_forward_function : callable
 383        Activation function used for dendrites.
 384    modules_to_perforate : list
 385        Module types to convert to PAI modules for perforation.
 386    module_names_to_perforate : list
 387        Module names to convert to PAI modules for perforation.
 388    module_ids_to_perforate : list
 389        Specific module IDs to convert to PAI modules for perforation.
 390    modules_to_track : list
 391        Module types to track but not convert.
 392    module_names_to_track : list
 393        Module names to track but not convert.
 394    module_ids_to_track : list
 395        Specific module IDs to track but not convert.
 396    modules_to_replace : list
 397        Module types to replace before conversion.
 398    replacement_modules : list
 399        Replacement modules for modules_to_replace.
 400    modules_with_processing : list
 401        Module types requiring custom processing.
 402    modules_processing_classes : list
 403        Processing classes for modules_with_processing.
 404    module_names_with_processing : list
 405        Module names requiring custom processing.
 406    module_by_name_processing_classes : list
 407        Processing classes for module_names_with_processing.
 408    module_names_to_not_save : list
 409        Module names to exclude from saving.
 410    perforated_backpropagation : bool
 411        Whether Perforated Backpropagation is enabled.
 412    """
 413
 414    # Explicit type map for every config variable — used by load_config to coerce JSON values.
 415    # [type] means "list whose elements are class/type objects" (need dotted-name resolution).
 416    _TYPES: dict = {
 417        **{
 418            k: bool
 419            for k in (
 420                "use_cuda",
 421                "confirm_correct_sizes",
 422                "unwrapped_modules_confirmed",
 423                "weight_decay_accepted",
 424                "checked_skipped_modules",
 425                "configuration_confirmed",
 426                "verbose",
 427                "extra_verbose",
 428                "silent",
 429                "save_old_graph_scores",
 430                "testing_dendrite_capacity",
 431                "strict_loading",
 432                "using_safe_tensors",
 433                "drawing_pai",
 434                "drawing_extra_graphs",
 435                "test_saves",
 436                "pai_saves",
 437                "reset_best_score_on_switch",
 438                "learn_dendrites_live",
 439                "no_extra_n_modes",
 440                "retain_all_dendrites",
 441                "find_best_lr",
 442                "dont_give_up_unless_learning_rate_lowered",
 443                "maximizing_score",
 444                "candidate_weight_init_by_main",
 445                "perforated_backpropagation",
 446                "weight_tying_experimental",
 447                "dashboard_events_enabled",
 448                "dashboard_debug",
 449            )
 450        },
 451        **{
 452            k: int
 453            for k in (
 454                "debugging_output_dimensions",
 455                "global_candidates",
 456                "n_epochs_to_switch",
 457                "history_lookback",
 458                "initial_history_after_switches",
 459                "fixed_switch_num",
 460                "first_fixed_switch_num",
 461                "switch_mode",
 462                "max_dendrite_tries",
 463                "max_dendrites",
 464                "param_vals_setting",
 465            )
 466        },
 467        **{
 468            k: float
 469            for k in (
 470                "improvement_threshold_raw",
 471                "candidate_weight_initialization_multiplier",
 472            )
 473        },
 474        **{
 475            k: str
 476            for k in (
 477                "save_name",
 478                "library_validation_score",
 479                "dashboard_url",
 480            )
 481        },
 482        "device": torch.device,
 483        "d_type": torch.dtype,
 484        "pai_forward_function": callable,
 485        **{
 486            k: list
 487            for k in (
 488                "output_dimensions",
 489                "improvement_threshold",
 490                "module_names_to_perforate",
 491                "module_ids_to_perforate",
 492                "module_names_to_track",
 493                "module_ids_to_track",
 494                "parameter_ids_to_track",
 495                "module_names_with_processing",
 496                "module_names_to_not_save",
 497                "library_extra_scores",
 498                "library_extra_scores_without_graphing",
 499            )
 500        },
 501        **{
 502            k: [type]
 503            for k in (
 504                "modules_to_perforate",
 505                "modules_to_track",
 506                "modules_to_replace",
 507                "replacement_modules",
 508                "modules_with_processing",
 509                "modules_processing_classes",
 510                "module_by_name_processing_classes",
 511            )
 512        },
 513    }
 514
 515    # Subset of _TYPES: variables that can be overridden on a per-module basis
 516    # via the Studio.  These are set outside the ``if not module_name:`` block
 517    # in :py:meth:`__init__`, so they are meaningful when constructing a
 518    # module-specific PAIConfig instance.
 519    _CUSTOMIZABLE: dict = {
 520        #"verbose": bool,
 521        #"extra_verbose": bool,
 522        #"silent": bool,
 523        #"global_candidates": int,
 524        "output_dimensions": list,
 525        #"candidate_weight_initialization_multiplier": float,
 526        #"candidate_weight_init_by_main": bool,
 527        #"retain_all_dendrites": bool,
 528        #"max_dendrites": int,
 529        #"pai_forward_function": callable,
 530    }
 531
 532    def __getstate__(self):
 533        """Tell pickle what to save when this object is serialized (e.g. torch.save).
 534
 535        The get_*/set_*/append_* methods attached to each PAIConfig instance are
 536        locally-defined closures and cannot be pickled.  Only the underlying data
 537        values (stored as _<name> in __dict__) need to be saved; __setstate__
 538        will recreate the methods by calling __init__ on load.
 539        """
 540        import types
 541
 542        # Walk every attribute on this instance and keep only the plain data values.
 543        # Bound methods (the get_*/set_*/append_* closures) are skipped because
 544        # they cannot be serialized by pickle.
 545        pickle_safe_state = {}
 546        for attr_name, attr_value in self.__dict__.items():
 547            if not isinstance(attr_value, types.MethodType):
 548                pickle_safe_state[attr_name] = attr_value
 549
 550        return pickle_safe_state
 551
 552    def __setstate__(self, saved_state):
 553        """Restore this object from a pickled state (e.g. torch.load).
 554
 555        Calls __init__ to rebuild all the get_*/set_*/append_* methods, then
 556        overlays the saved data values on top so all settings are preserved.
 557        """
 558        # Grab the identity fields needed to reconstruct the right flavor of config
 559        # (global config vs. per-module config).
 560        saved_module_name = saved_state.get("_module_name")
 561        saved_module_type = saved_state.get("_module_type")
 562
 563        # Re-run __init__ so all the get_*/set_*/append_* methods exist again.
 564        # This also sets default values for every variable, which we overwrite below.
 565        self.__init__(module_name=saved_module_name, module_type=saved_module_type)
 566
 567        # Overwrite the defaults with the actual saved values.
 568        # Only restore private data attributes (prefixed with _); skip the identity
 569        # fields we already handled above, and skip _config_file for now so that
 570        # the setter calls above do not trigger auto-save during the restore.
 571        skip_on_first_pass = {"_module_name", "_module_type", "_config_file"}
 572        for attr_name, attr_value in saved_state.items():
 573            if attr_name.startswith("_") and attr_name not in skip_on_first_pass:
 574                self.__dict__[attr_name] = attr_value
 575
 576        # Restore the config file path last.  Auto-save is driven by _config_file
 577        # being set, so we only put it back after all values are already in place.
 578        if saved_state.get("_config_file"):
 579            self.__dict__["_config_file"] = saved_state["_config_file"]
 580
 581    def __getattr__(self, name):
 582        """Handle missing attributes gracefully, especially for PB variables.
 583
 584        Parameters
 585        ----------
 586        name : str
 587            The name of the attribute being accessed.
 588
 589        Returns
 590        -------
 591        None or raises AttributeError
 592            Returns None for missing set_ methods, raises AttributeError otherwise.
 593        """
 594        if name.startswith("set_"):
 595            print(f"Variable '{name[4:]}' does not exist.  Ignoring set attempt.")
 596            return lambda x: None
 597        if name.startswith("append_"):
 598            print(
 599                f"List Variable '{name[7:]}' does not exist.  Ignoring append attempt."
 600            )
 601            return lambda x: None
 602        if name.startswith("get_") and self.__dict__.get("_module_name") is not None:
 603            # Module-specific config: check for a per-module override stored directly in
 604            # __dict__ (written by load_config when custom JSON data was found for this
 605            # module).  This covers CUSTOMIZABLE vars that are not initialised for
 606            # per-module configs (e.g. output_dimensions, which lives inside the
 607            # ``if not module_name:`` block).
 608            private_key = f"_{name[4:]}"
 609            if private_key in self.__dict__:
 610                stored = self.__dict__[private_key]
 611                return lambda: stored
 612            # Fall back to the global pc instance for vars not set on this instance.
 613            global_getter = getattr(pc, name, None)
 614            if global_getter is not None:
 615                return global_getter
 616        raise AttributeError(
 617            f"'{self.__class__.__name__}' object has no attribute '{name}'"
 618        )
 619
 620    def __init__(self, module_name=None, module_type=None):
 621        """Initialize PAIConfig with default settings.
 622
 623        module_name=None means this is the main global config.
 624        If module_name is set this is a per-module config that loads
 625        custom settings from module_settings[module_name] (by id) or
 626        module_settings[module_type] (by type) in the JSON file.
 627        """
 628        # Must be first: prevents __getattr__ from firing for _config_file
 629        # during construction (before add_pai_config_var_functions sets it).
 630        # Also disables auto-save in setters until the end of __init__.
 631        self.__dict__["_config_file"] = None
 632        self.__dict__["_manually_set_keys"] = set()
 633        self.__dict__["_loading_config_values"] = False
 634        self.__dict__["_auto_persist_config"] = True
 635        # None = global config; any string = per-module config
 636        self.__dict__["_module_name"] = module_name
 637        # Short class name of the wrapped module (e.g. 'Conv2d'), used as
 638        # a fallback lookup key when the specific name has no saved settings.
 639        self.__dict__["_module_type"] = module_type
 640
 641        if not module_name:
 642            ### Global Constants
 643            # Device configuration
 644            self.use_cuda = torch.cuda.is_available()
 645            add_pai_config_var_functions(self, "use_cuda", self.use_cuda)
 646            self.device = torch.device("cuda" if self.use_cuda else "cpu")
 647            add_pai_config_var_functions(self, "device", self.device)
 648
 649            self.save_name = ""
 650            add_pai_config_var_functions(self, "save_name", self.save_name)
 651
 652            self.config_file = None
 653            add_pai_config_var_functions(self, "config_file", self.config_file)
 654
 655            # Debug settings
 656            self.debugging_output_dimensions = 0
 657            add_pai_config_var_functions(
 658                self, "debugging_output_dimensions", self.debugging_output_dimensions
 659            )
 660            # Debugging input tensor sizes.
 661            # This will slow things down very slightly and is not necessary but can help
 662            # catch when dimensions were not filled in correctly.
 663            self.confirm_correct_sizes = False
 664            add_pai_config_var_functions(
 665                self, "confirm_correct_sizes", self.confirm_correct_sizes
 666            )
 667
 668            # Confirmation flags for non-recommended options
 669            self.unwrapped_modules_confirmed = False
 670            add_pai_config_var_functions(
 671                self, "unwrapped_modules_confirmed", self.unwrapped_modules_confirmed
 672            )
 673            # We used to be worried this caused problems, but it is now accepted.
 674            # Leaving code in case future changes require revisiting this decision.
 675            self.weight_decay_accepted = True
 676            add_pai_config_var_functions(
 677                self, "weight_decay_accepted", self.weight_decay_accepted
 678            )
 679            self.checked_skipped_modules = False
 680            add_pai_config_var_functions(
 681                self, "checked_skipped_modules", self.checked_skipped_modules
 682            )
 683            self.configuration_confirmed = False
 684            add_pai_config_var_functions(
 685                self, "configuration_confirmed", self.configuration_confirmed
 686            )
 687            # Analysis settings
 688            self.save_old_graph_scores = True
 689            add_pai_config_var_functions(
 690                self, "save_old_graph_scores", self.save_old_graph_scores
 691            )
 692            # Testing settings
 693            self.testing_dendrite_capacity = True
 694            add_pai_config_var_functions(
 695                self, "testing_dendrite_capacity", self.testing_dendrite_capacity
 696            )
 697
 698            # File format settings
 699            self.using_safe_tensors = True
 700            add_pai_config_var_functions(
 701                self, "using_safe_tensors", self.using_safe_tensors
 702            )
 703
 704            # Checkpoint loading settings
 705            # Whether to use strict=True when loading state_dict
 706            # Set to False if loading old checkpoints that are missing new fields
 707            self.strict_loading = True
 708            add_pai_config_var_functions(
 709                self, "strict_loading", self.strict_loading
 710            )
 711
 712            # Graph and visualization settings
 713            # A graph setting which can be set to false if you want to do your own
 714            # training visualizations
 715            self.drawing_pai = True
 716            add_pai_config_var_functions(self, "drawing_pai", self.drawing_pai)
 717
 718            # Drawing extra graphs beyond the standard ones.
 719            self.drawing_extra_graphs = True
 720            add_pai_config_var_functions(
 721                self, "drawing_extra_graphs", self.drawing_extra_graphs
 722            )
 723
 724            # Saving test intermediary models, good for experimentation, bad for memory
 725            self.test_saves = True
 726            add_pai_config_var_functions(self, "test_saves", self.test_saves)
 727            # To be filled in later. pai_saves will remove some extra scaffolding for
 728            # slight memory and speed improvements
 729            self.pai_saves = False
 730            add_pai_config_var_functions(self, "pai_saves", self.pai_saves)
 731            # Improvement thresholds
 732            # Percentage improvement increase needed to call a new best validation score
 733            self.improvement_threshold = [0.001, 0.0001, 0.0]
 734            add_pai_config_var_functions(
 735                self, "improvement_threshold", self.improvement_threshold
 736            )
 737
 738            # Raw increase needed
 739            self.improvement_threshold_raw = 1e-5
 740            add_pai_config_var_functions(
 741                self, "improvement_threshold_raw", self.improvement_threshold_raw
 742            )
 743            # SWITCH MODE SETTINGS
 744
 745            # Add dendrites every time to debug implementation
 746            self.DOING_SWITCH_EVERY_TIME = 0
 747
 748            # Switch when validation hasn't improved over x epochs
 749            self.DOING_HISTORY = 1
 750            # Epochs to try before deciding to load previous best and add dendrites
 751            # Be sure this is higher than scheduler patience
 752            self.n_epochs_to_switch = 10
 753            add_pai_config_var_functions(
 754                self, "n_epochs_to_switch", self.n_epochs_to_switch
 755            )
 756            # Number to average validation scores over
 757            self.history_lookback = 1
 758            add_pai_config_var_functions(
 759                self, "history_lookback", self.history_lookback
 760            )
 761            # Amount of epochs to run after adding a new set of dendrites before checking
 762            # to add more
 763            self.initial_history_after_switches = 0
 764            add_pai_config_var_functions(
 765                self,
 766                "initial_history_after_switches",
 767                self.initial_history_after_switches,
 768            )
 769
 770            # Switch after a fixed number of epochs
 771            self.DOING_FIXED_SWITCH = 2
 772            # Number of epochs to complete before switching
 773            self.fixed_switch_num = 250
 774            add_pai_config_var_functions(
 775                self, "fixed_switch_num", self.fixed_switch_num
 776            )
 777            # An additional flag if you want your first switch to occur later than all the
 778            # rest for initial pretraining.  This is a new minimum, if its lower than
 779            # the above it will be ignored.
 780            self.first_fixed_switch_num = -1
 781            add_pai_config_var_functions(
 782                self, "first_fixed_switch_num", self.first_fixed_switch_num
 783            )
 784
 785            # A setting to not add dendrites and just do regular training
 786            # Warning, this will also never trigger training_complete
 787            self.DOING_NO_SWITCH = 3
 788
 789            # Default switch mode
 790            self.switch_mode = self.DOING_HISTORY
 791            add_pai_config_var_functions(self, "switch_mode", self.switch_mode)
 792
 793            # Reset settings
 794            # Resets score on switch
 795            # This can be useful if you need many epochs to catch up to the best score
 796            # from the previous version after adding dendrites
 797            self.reset_best_score_on_switch = True
 798            add_pai_config_var_functions(
 799                self, "reset_best_score_on_switch", self.reset_best_score_on_switch
 800            )
 801
 802            # Advanced settings
 803            # Not used in open source implementation, leave as default
 804            self.learn_dendrites_live = False
 805            add_pai_config_var_functions(
 806                self, "learn_dendrites_live", self.learn_dendrites_live
 807            )
 808            self.no_extra_n_modes = True
 809            add_pai_config_var_functions(
 810                self, "no_extra_n_modes", self.no_extra_n_modes
 811            )
 812
 813            # Data type for new modules and dendrite to dendrite / dendrite to neuron
 814            # weights
 815            self.d_type = torch.float
 816            add_pai_config_var_functions(self, "d_type", self.d_type)
 817
 818            # Learning rate management
 819            # A setting to automatically sweep over previously used learning rates when
 820            # adding new dendrites
 821            # Sometimes it's best to go back to initial LR, but often its best to start
 822            # at a lower LR
 823            self.find_best_lr = True
 824            add_pai_config_var_functions(self, "find_best_lr", self.find_best_lr)
 825            # Enforces the above even if the previous epoch didn't lower the learning rate
 826            self.dont_give_up_unless_learning_rate_lowered = True
 827            add_pai_config_var_functions(
 828                self,
 829                "dont_give_up_unless_learning_rate_lowered",
 830                self.dont_give_up_unless_learning_rate_lowered,
 831            )
 832            # Whether a higher validation score is better (True) or lower is better (False).
 833            self.maximizing_score = True
 834            add_pai_config_var_functions(
 835                self, "maximizing_score", self.maximizing_score
 836            )
 837
 838            # Dendrite attempt settings
 839            # Set to 1 if you want to quit as soon as one dendrite fails
 840            # Higher values will try new random dendrite weights this many times before
 841            # accepting that more dendrites don't improve
 842            self.max_dendrite_tries = 2
 843            add_pai_config_var_functions(
 844                self, "max_dendrite_tries", self.max_dendrite_tries
 845            )
 846
 847            # Scheduler parameter settings
 848            # Have learning rate params be by total epoch
 849            self.PARAM_VALS_BY_TOTAL_EPOCH = 0
 850            # Reset the params at every switch
 851            self.PARAM_VALS_BY_UPDATE_EPOCH = 1
 852            # Reset params for dendrite starts but not for normal restarts
 853            # Not used for open source version
 854            self.PARAM_VALS_BY_NEURON_EPOCH_START = 2
 855            # Default setting
 856            self.param_vals_setting = self.PARAM_VALS_BY_UPDATE_EPOCH
 857            add_pai_config_var_functions(
 858                self, "param_vals_setting", self.param_vals_setting
 859            )
 860            # Lists for module types and names to add dendrites to
 861            # For these lists no specifier means type, name is module name
 862            # and ids is the individual modules id, eg. model.conv2
 863            self.modules_to_perforate = []
 864            add_pai_config_var_functions(
 865                self, "modules_to_perforate", self.modules_to_perforate, list_type=True
 866            )
 867            self.module_names_to_perforate = [
 868                "PAISequential",
 869                "Conv1d",
 870                "Conv2d",
 871                "Conv3d",
 872                "Linear",
 873            ]
 874            add_pai_config_var_functions(
 875                self,
 876                "module_names_to_perforate",
 877                self.module_names_to_perforate,
 878                list_type=True,
 879            )
 880            self.module_ids_to_perforate = []
 881            add_pai_config_var_functions(
 882                self,
 883                "module_ids_to_perforate",
 884                self.module_ids_to_perforate,
 885                list_type=True,
 886            )
 887
 888            # All modules should either be perforated or tracked to ensure all modules
 889            # are accounted for
 890            self.modules_to_track = []
 891            add_pai_config_var_functions(
 892                self, "modules_to_track", self.modules_to_track, list_type=True
 893            )
 894            self.module_names_to_track = []
 895            add_pai_config_var_functions(
 896                self,
 897                "module_names_to_track",
 898                self.module_names_to_track,
 899                list_type=True,
 900            )
 901            # IDs are for if you want to pass only a single module by its assigned ID rather than the module type by name
 902            self.module_ids_to_track = []
 903            add_pai_config_var_functions(
 904                self, "module_ids_to_track", self.module_ids_to_track, list_type=True
 905            )
 906
 907            # Parameter IDs to track as neuron parameters without recursive behavior
 908            # (e.g., [".my_custom_parameter"]).
 909            self.parameter_ids_to_track = []
 910            add_pai_config_var_functions(
 911                self,
 912                "parameter_ids_to_track",
 913                self.parameter_ids_to_track,
 914                list_type=True,
 915            )
 916
 917            # Replacement modules happen before the conversion,
 918            # so replaced modules will then also be run through the conversion steps
 919            # These are for modules that need to be replaced before addition of dendrites
 920            # See the resnet example in models_perforatedai
 921            self.modules_to_replace = []
 922            add_pai_config_var_functions(
 923                self, "modules_to_replace", self.modules_to_replace, list_type=True
 924            )
 925            # Modules to replace the above modules with
 926            self.replacement_modules = []
 927            add_pai_config_var_functions(
 928                self, "replacement_modules", self.replacement_modules, list_type=True
 929            )
 930
 931            # Dendrites default to modules which are one tensor input and one tensor
 932            # output in forward()
 933            # Other modules require to be labeled as modules with processing and assigned
 934            # processing classes
 935            # This can be done by module type or module name see customization.md in API
 936            # for example
 937            self.modules_with_processing = []
 938            add_pai_config_var_functions(
 939                self,
 940                "modules_with_processing",
 941                self.modules_with_processing,
 942                list_type=True,
 943            )
 944            self.modules_processing_classes = []
 945            add_pai_config_var_functions(
 946                self,
 947                "modules_processing_classes",
 948                self.modules_processing_classes,
 949                list_type=True,
 950            )
 951            self.module_names_with_processing = []
 952            add_pai_config_var_functions(
 953                self,
 954                "module_names_with_processing",
 955                self.module_names_with_processing,
 956                list_type=True,
 957            )
 958            self.module_by_name_processing_classes = []
 959            add_pai_config_var_functions(
 960                self,
 961                "module_by_name_processing_classes",
 962                self.module_by_name_processing_classes,
 963                list_type=True,
 964            )
 965
 966            # Similarly here as above. Some huggingface models have multiple pointers to
 967            # the same modules which cause problems
 968            # If you want to only save one of the multiple pointers you can set which ones
 969            # not to save here
 970            self.module_names_to_not_save = [".base_model"]
 971            add_pai_config_var_functions(
 972                self,
 973                "module_names_to_not_save",
 974                self.module_names_to_not_save,
 975                list_type=True,
 976            )
 977
 978            # Perforated Backpropagation settings
 979            self.perforated_backpropagation = False
 980            add_pai_config_var_functions(
 981                self, "perforated_backpropagation", self.perforated_backpropagation
 982            )
 983
 984            # This is specifically a workaround for weight tying
 985            # Not to be used for a duplicate pointer that isn't actually run twice
 986            self.weight_tying_experimental = False
 987            add_pai_config_var_functions(
 988                self, "weight_tying_experimental", self.weight_tying_experimental
 989            )
 990
 991            # Dashboard event streaming settings
 992            self.dashboard_events_enabled = False
 993            add_pai_config_var_functions(
 994                self, "dashboard_events_enabled", self.dashboard_events_enabled
 995            )
 996            self.dashboard_url = "http://localhost:3002"
 997            add_pai_config_var_functions(
 998                self, "dashboard_url", self.dashboard_url
 999            )
1000            self.dashboard_debug = False
1001            add_pai_config_var_functions(
1002                self, "dashboard_debug", self.dashboard_debug
1003            )
1004
1005            # These are settings where libraries must be doing the scoring adding to
1006            # message from your main script what metric to use
1007            self.library_validation_score = ""
1008            add_pai_config_var_functions(
1009                self, "library_validation_score", self.library_validation_score
1010            )
1011            self.library_extra_scores = []
1012            add_pai_config_var_functions(
1013                self,
1014                "library_extra_scores",
1015                self.library_extra_scores,
1016                list_type=True,
1017            )
1018            self.library_extra_scores_without_graphing = []
1019            add_pai_config_var_functions(
1020                self,
1021                "library_extra_scores_without_graphing",
1022                self.library_extra_scores_without_graphing,
1023                list_type=True,
1024            )
1025
1026            # Input dimensions needs to be set every time. It is set to what format of
1027            # planes you are expecting.
1028            # Neuron index should be set to 0, variable indexes should be set to -1.
1029            # For example, if your format is [batchsize, nodes, x, y]
1030            # output_dimensions is [-1, 0, -1, -1].
1031            # if your format is, [batchsize, time index, nodes] output_dimensions is
1032            # [-1, -1, 0]
1033            self.output_dimensions = [-1, 0, -1, -1]
1034            add_pai_config_var_functions(
1035                self, "output_dimensions", self.output_dimensions, list_type=True
1036            )
1037        # Verbosity settings
1038        self.verbose = False
1039        add_pai_config_var_functions(self, "verbose", self.verbose)
1040        self.extra_verbose = False
1041        add_pai_config_var_functions(self, "extra_verbose", self.extra_verbose)
1042        # Suppress all PAI prints
1043        self.silent = False
1044        add_pai_config_var_functions(self, "silent", self.silent)
1045
1046        # In place for future implementation options of adding multiple candidate
1047        # dendrites together
1048        self.global_candidates = 1
1049        add_pai_config_var_functions(self, "global_candidates", self.global_candidates)
1050
1051        # Weight initialization settings
1052        # Multiplier when randomizing dendrite weights
1053        self.candidate_weight_initialization_multiplier = 0.01
1054        add_pai_config_var_functions(
1055            self,
1056            "candidate_weight_initialization_multiplier",
1057            self.candidate_weight_initialization_multiplier,
1058        )
1059        # Multiplier when randomizing dendrite weights
1060        self.candidate_weight_init_by_main = False
1061        add_pai_config_var_functions(
1062            self,
1063            "candidate_weight_init_by_main",
1064            self.candidate_weight_init_by_main,
1065        )
1066
1067        # Dendrite retention settings
1068        # A setting to keep dendrites even if they do not improve scores
1069        self.retain_all_dendrites = False
1070        add_pai_config_var_functions(
1071            self, "retain_all_dendrites", self.retain_all_dendrites
1072        )
1073
1074        # Max dendrites to add even if they do continue improving scores
1075        self.max_dendrites = 100
1076        add_pai_config_var_functions(self, "max_dendrites", self.max_dendrites)
1077
1078        # Activation function settings
1079        # The activation function to use for dendrites
1080        self.pai_forward_function = torch.sigmoid
1081        add_pai_config_var_functions(
1082            self, "pai_forward_function", self.pai_forward_function
1083        )
1084
1085        # Apply per-module overrides after defaults are defined so loaded values
1086        # win over initializer defaults.
1087        if module_name is not None:
1088            import os
1089
1090            global_pc = globals().get("pc")
1091            source_paths = []
1092            if global_pc is not None:
1093                config_file = getattr(global_pc, "__dict__", {}).get("_config_file")
1094                run_config = global_pc.get_run_config_path()
1095                if config_file:
1096                    source_paths.append(config_file)
1097                if run_config:
1098                    source_paths.append(run_config)
1099
1100            for source_path in source_paths:
1101                if source_path and os.path.exists(source_path):
1102                    self.load_config(
1103                        source_path,
1104                        module_name=module_name,
1105                        module_type=module_type,
1106                    )
1107                    break
1108
1109        # ------------------------------------------------------------------
1110        # Config file will be set when save_name is assigned (in perforate_model)
1111        # ------------------------------------------------------------------
1112        # _config_file stays None until save_name is set to a non-empty value
1113
1114    # ------------------------------------------------------------------
1115
1116    def save_config(self, filename):
1117        """Save the current PAIConfig state to a JSON file.
1118
1119        Parameters
1120        ----------
1121        filename : str
1122            Destination file path (created or overwritten).
1123
1124        Notes
1125        -----
1126        Values that are not natively JSON-serialisable (torch.device,
1127        torch.dtype, nn.Module subclasses, callables) are stored as their
1128        dotted-string representations so they can be round-tripped by
1129        :py:meth:`load_config`.
1130
1131        Returns
1132        -------
1133        None
1134            This function does not return a value.
1135        """
1136        import json
1137
1138        config_dict = {}
1139
1140        # Private-storage vars added by add_pai_config_var_functions
1141        # e.g. self._module_ids_to_perforate → key 'module_ids_to_perforate'
1142        for key, val in sorted(self.__dict__.items()):
1143            if not (key.startswith("_") and not key.startswith("__")):
1144                continue
1145            # Skip internal bookkeeping keys that must not round-trip through JSON
1146            if key in (
1147                "_config_file",
1148                "_module_name",
1149                "_module_type",
1150                "_manually_set_keys",
1151                "_loading_config_values",
1152                "_auto_persist_config",
1153            ):
1154                continue
1155            clean_key = key[1:]
1156            # Keep callable config fields (e.g. pai_forward_function) but skip
1157            # any other private callable values that are not typed as callable.
1158            if callable(val) and PAIConfig._TYPES.get(clean_key) is not callable:
1159                continue
1160            try:
1161                config_dict[clean_key] = _serialize_pai_value(val)
1162            except Exception:
1163                config_dict[clean_key] = str(val)
1164
1165        # Plain constants (DOING_*, PARAM_VALS_BY_*, etc.)
1166        for key, val in self.__dict__.items():
1167            if key.startswith("_") or callable(val):
1168                continue
1169            if key not in config_dict:
1170                # Keep config_file runtime-only: do not persist pointer paths in
1171                # saved JSON snapshots.
1172                if key == "config_file":
1173                    continue
1174                try:
1175                    config_dict[key] = _serialize_pai_value(val)
1176                except Exception:
1177                    config_dict[key] = str(val)
1178
1179        # Merge short class names from modules_to_perforate into module_names_to_perforate
1180        # so the UI (and JS) only needs to check one array.
1181        type_short_names = [
1182            cls.__name__
1183            for cls in self.__dict__.get("_modules_to_perforate", [])
1184            if isinstance(cls, type)
1185        ]
1186        existing = config_dict.get("module_names_to_perforate", [])
1187        config_dict["module_names_to_perforate"] = existing + [
1188            n for n in type_short_names if n not in existing
1189        ]
1190
1191        # Preserve any per-module settings written by the Studio frontend.
1192        _existing_ms: dict = {}
1193        try:
1194            with open(filename, "r") as _f:
1195                _existing_ms = json.load(_f).get("module_settings", {})
1196        except Exception:
1197            pass
1198        config_dict["module_settings"] = _existing_ms
1199
1200        # Publish customizable field type names so the Studio frontend can
1201        # render appropriate editors without needing to import the library.
1202        config_dict["_customizable_fields"] = {
1203            k: (v.__name__ if hasattr(v, "__name__") else str(v))
1204            for k, v in PAIConfig._CUSTOMIZABLE.items()
1205        }
1206
1207        # Ensure the directory exists before saving
1208        import os
1209
1210        directory = os.path.dirname(filename)
1211        if directory:
1212            os.makedirs(directory, exist_ok=True)
1213
1214        with open(filename, "w") as f:
1215            json.dump(config_dict, f, indent=2)
1216        print(f"[PAI Config] Saved {len(config_dict)} variables \u2192 {filename}")
1217
1218    def load_config(self, filename, module_name=None, module_type=None, skip_keys=None):
1219        """Load PAIConfig state from a JSON file produced by :py:meth:`save_config`.
1220
1221        If *module_name* is ``None`` (default) every serialisable variable in
1222        the file is restored on this instance.
1223
1224        If *module_name* is given the lookup priority is:
1225          1. ``module_settings[module_name]`` (exact name / id match)
1226          2. ``module_settings[module_type]``  (type-level fallback)
1227          3. No-op — the defaults already set by ``__init__`` are kept.
1228
1229        Parameters
1230        ----------
1231        filename : str
1232            Path to the JSON file to read.
1233        module_name : str, optional
1234            Display name (id) of the module whose custom settings should be loaded.
1235        module_type : str, optional
1236            Short class name of the module type, used as a fallback key.
1237
1238        skip_keys : set, optional
1239            Variable names to ignore while loading (manual overrides).
1240
1241        Returns
1242        -------
1243        None
1244            This function does not return a value.
1245        """
1246        import json
1247
1248        with open(filename, "r") as f:
1249            config_dict = json.load(f)
1250
1251        if module_name is not None:
1252            # ── Per-module load ──────────────────────────────────────────────
1253            module_settings = config_dict.get("module_settings", {})
1254            # Priority: exact name → type fallback → no-op
1255            if module_name in module_settings:
1256                custom = module_settings[module_name]
1257                resolved_key = module_name
1258            elif module_type and module_type in module_settings:
1259                custom = module_settings[module_type]
1260                resolved_key = module_type
1261            else:
1262                # No custom settings for this module or type — keep defaults.
1263                return
1264            loaded = 0
1265            skipped = 0
1266            for key, json_val in custom.items():
1267                if key not in PAIConfig._CUSTOMIZABLE:
1268                    continue
1269                type_hint = PAIConfig._TYPES.get(key)
1270                private_key = f"_{key}"
1271                # Write directly to __dict__ so we bypass any setter guards and also
1272                # correctly handle vars that are not pre-initialised on per-module
1273                # configs (e.g. output_dimensions, which only lives inside the
1274                # ``if not module_name:`` block of __init__).
1275                try:
1276                    self.__dict__[private_key] = (
1277                        _deserialize_pai_value(json_val, type_hint)
1278                        if type_hint is not None
1279                        else json_val
1280                    )
1281                    loaded += 1
1282                except Exception as exc:
1283                    print(
1284                        f"[PAI Config] Warning: could not load '{key}' for '{resolved_key}': {exc}"
1285                    )
1286                    skipped += 1
1287            print(
1288                f"[PAI Config] Loaded {loaded} custom vars for '{resolved_key}' from {filename}"
1289                + (f" ({skipped} skipped)" if skipped else "")
1290            )
1291            return
1292
1293        # ── Global load: every variable in the JSON ──────────────────────────
1294        skip_keys = set(skip_keys or [])
1295        loaded = 0
1296        skipped = 0
1297        self.__dict__["_loading_config_values"] = True
1298        try:
1299            for key, json_val in config_dict.items():
1300                # Skip internal bookkeeping and Studio-only metadata keys.
1301                # 'module_name' and 'module_type' must never overwrite the
1302                # instance's _module_name/_module_type (they are internal only).
1303                if key in (
1304                    "config_file",
1305                    "module_settings",
1306                    "module_name",
1307                    "module_type",
1308                ) or key.startswith("_"):
1309                    continue
1310                if key in skip_keys:
1311                    continue
1312                type_hint = PAIConfig._TYPES.get(key)
1313                private_key = f"_{key}"
1314                if hasattr(self, private_key):
1315                    try:
1316                        setattr(
1317                            self,
1318                            private_key,
1319                            (
1320                                _deserialize_pai_value(json_val, type_hint)
1321                                if type_hint is not None
1322                                else json_val
1323                            ),
1324                        )
1325                        loaded += 1
1326                    except Exception as exc:
1327                        print(f"[PAI Config] Warning: could not load '{key}': {exc}")
1328                        skipped += 1
1329                elif hasattr(self, key) and not callable(getattr(self, key, None)):
1330                    try:
1331                        setattr(self, key, json_val)
1332                        loaded += 1
1333                    except Exception:
1334                        skipped += 1
1335        finally:
1336            self.__dict__["_loading_config_values"] = False
1337
1338        print(
1339            f"[PAI Config] Loaded {loaded} variables from {filename}"
1340            + (f" ({skipped} skipped)" if skipped else "")
1341        )
1342
1343    def get_run_config_path(self):
1344        """Get the run-specific config path derived from save_name.
1345
1346        Returns
1347        -------
1348        str or None
1349            Path to the run config file, or None when save_name is unset.
1350        """
1351        import os
1352
1353        save_name = self.__dict__.get("_save_name")
1354        if not save_name:
1355            return None
1356        return os.path.join(os.getcwd(), save_name, f"{save_name}_config.json")
1357
1358    def sync_config_sources(self):
1359        """Load config values from configured local JSON sources.
1360
1361        Load order:
1362        1. config_file pointer (if set)
1363        2. run config derived from save_name (if present)
1364
1365        Keys set manually in this process are never overwritten.
1366        """
1367        import os
1368
1369        if self.__dict__.get("_module_name") is not None:
1370            return
1371
1372        skip_keys = self.__dict__.get("_manually_set_keys", set())
1373        config_file = self.__dict__.get("_config_file")
1374        run_config = self.get_run_config_path()
1375
1376        if config_file and os.path.exists(config_file):
1377            self.load_config(config_file, skip_keys=skip_keys)
1378        if run_config and os.path.exists(run_config):
1379            self.load_config(run_config, skip_keys=skip_keys)
1380
1381    def persist_config_outputs(self, overwrite_config_file=False):
1382        """Persist current config state to local JSON outputs.
1383
1384        Parameters
1385        ----------
1386        overwrite_config_file : bool
1387            When True, also writes the configured config_file pointer target.
1388        """
1389        if self.__dict__.get("_module_name") is not None:
1390            return
1391
1392        run_config = self.get_run_config_path()
1393        if run_config:
1394            self.save_config(run_config)
1395
1396        if overwrite_config_file:
1397            config_file = self.__dict__.get("_config_file")
1398            if config_file:
1399                self.save_config(config_file)
1400
1401
1402class PAISequential(nn.Sequential):
1403    """Sequential module wrapper for PAI.
1404
1405    This wrapper takes an array of layers and creates a sequential container
1406    that is compatible with PAI's dendrite addition system. It should be used
1407    for normalization layers and can be used for final output layers.
1408
1409    Parameters
1410    ----------
1411    layer_array : list
1412        List of PyTorch nn.Module objects to be executed sequentially.
1413
1414    Examples
1415    --------
1416    >>> layers = [nn.Linear(2 * hidden_dim, seq_width),
1417    ...           nn.LayerNorm(seq_width)]
1418    >>> sequential_block = PAISequential(layers)
1419
1420    Notes
1421    -----
1422    This should be used for:
1423        - All normalization layers (LayerNorm, BatchNorm, etc.)
1424    This can be used for:
1425        - Final output layer and softmax combinations
1426    """
1427
1428    def __init__(self, layer_array):
1429        """Initialize PAISequential with a list of layers.
1430
1431        Parameters
1432        ----------
1433        layer_array : list
1434            List of PyTorch modules to execute in sequence.
1435        """
1436        super(PAISequential, self).__init__()
1437        self.model = nn.Sequential(*layer_array)
1438
1439    def forward(self, *args, **kwargs):
1440        """Forward pass through the sequential layers.
1441
1442        Parameters
1443        ----------
1444        *args
1445            Positional arguments passed to the first layer.
1446        **kwargs
1447            Keyword arguments passed to the layers.
1448
1449        Returns
1450        -------
1451        torch.Tensor
1452            Output from the final layer in the sequence.
1453        """
1454        return self.model(*args, **kwargs)
1455
1456
1457### Global objects and variables
1458
1459### Global Modules
1460pc = PAIConfig()
1461"""Global PAIConfig instance.
1462
1463This is the primary configuration object used throughout the PAI system.
1464Modify settings through this instance to control PAI behavior.
1465"""
1466
1467"""Pointer to the PAI Tracker.
1468
1469This will be populated with the PAI Tracker instance which handles
1470the addition of dendrites during training. Initially an empty list.
1471"""
1472pai_tracker = []
1473
1474pai_scaler = None
1475
1476# This will be set to true if perforated backpropagation is available
1477# Do not just set this to True without the library and a license, it will cause errors
1478try:
1479    import perforatedbp.globals_pbp as perforatedbp_globals
1480
1481    print("Building dendrites with Perforated Backpropagation")
1482
1483    pc.set_perforated_backpropagation(True)
1484    # This is default to True for open source version
1485    # But defaults to False for perforated backpropagation
1486    pc.set_no_extra_n_modes(False)
1487
1488    # Loop through the vars module's attributes and add them dynamically
1489    for var_name in dir(perforatedbp_globals):
1490        if not var_name.startswith("_"):
1491            add_pai_config_var_functions(
1492                pc, var_name, getattr(perforatedbp_globals, var_name)
1493            )
1494
1495    # Merge PBP type hints into PAIConfig._TYPES so load_config can correctly
1496    # round-trip all perforatedbp variables from JSON.
1497    if hasattr(perforatedbp_globals, "_TYPES"):
1498        PAIConfig._TYPES.update(perforatedbp_globals._TYPES)
1499
1500except ImportError:
1501    print("Building dendrites without Perforated Backpropagation")