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        # Auto-save: if a config file has been configured (set when save_name is set),
 144        # persist the new value immediately so the JSON stays in sync.
 145        config_file = self.__dict__.get("_config_file")
 146        # Special case: if save_name changed to non-empty, update config file path
 147        if var_name == "save_name" and value:
 148            import os as _os
 149
 150            _save_folder = _os.path.join(_os.getcwd(), value)
 151            config_file = _os.path.join(_save_folder, f"{value}_config.json")
 152            self.__dict__["_config_file"] = config_file
 153        elif config_file and not self.__dict__.get("_testing_dendrite_capacity", False):
 154            self.save_config(config_file)
 155
 156    def appender(self, value):
 157        """Append a value to the property if it is a list.
 158
 159        Parameters
 160        ----------
 161        value : Any
 162            Value or values to append to the underlying list property.
 163
 164        Returns
 165        -------
 166        None
 167            Appends items in-place and prints the resulting list value.
 168        """
 169        if isinstance(getattr(self, private_name), list):
 170            if var_name in (
 171                "module_ids_to_track",
 172                "module_ids_to_perforate",
 173                "parameter_ids_to_track",
 174            ):
 175                for module_id in value:
 176                    _validate_module_id(module_id)
 177            setattr(self, private_name, getattr(self, private_name) + value)
 178            print(
 179                'New list value of "{}": {}'.format(
 180                    private_name, getattr(self, private_name)
 181                )
 182            )
 183        else:
 184            raise TypeError(f"Cannot append to non-list attribute '{var_name}'")
 185
 186    # Attach methods to the instance
 187    if list_type:
 188        setattr(obj, f"get_{var_name}", getter_list.__get__(obj))
 189    else:
 190        setattr(obj, f"get_{var_name}", getter_val.__get__(obj))
 191    setattr(obj, f"set_{var_name}", setter.__get__(obj))
 192    setattr(obj, f"append_{var_name}", appender.__get__(obj))
 193
 194
 195# ---------------------------------------------------------------------------
 196# JSON serialization helpers  (used by PAIConfig.save_config / load_config)
 197# ---------------------------------------------------------------------------
 198
 199
 200def _resolve_dotted_name(dotted_name):
 201    """Import and return an object identified by a dotted module path.
 202
 203    E.g. 'torch.nn.modules.conv.Conv2d'  → nn.Conv2d class
 204         'torch.sigmoid'                 → torch.sigmoid function
 205
 206    Returns None if the name cannot be resolved.
 207    """
 208    import importlib
 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
 223def _serialize_pai_value(val):
 224    """Recursively convert a PAIConfig value to a JSON-serialisable form."""
 225    if isinstance(val, bool):
 226        return val
 227    if isinstance(val, (int, float, str)):
 228        return val
 229    if isinstance(val, torch.device):
 230        return str(val)
 231    if isinstance(val, torch.dtype):
 232        return str(val)
 233    if isinstance(val, list):
 234        return [_serialize_pai_value(v) for v in val]
 235    if isinstance(val, type):
 236        mod = getattr(val, "__module__", "") or ""
 237        return f"{mod}.{val.__name__}" if mod else val.__name__
 238    if callable(val):
 239        name = getattr(val, "__name__", None)
 240        mod = getattr(val, "__module__", None)
 241        if name and mod:
 242            return f"{mod}.{name}"
 243        return str(val)
 244    return str(val)
 245
 246
 247def _deserialize_pai_value(json_val, type_hint):
 248    """Convert a JSON value to its Python type using an explicit type hint."""
 249    if json_val is None:
 250        return None
 251    if type_hint is torch.device:
 252        return torch.device(str(json_val))
 253    if type_hint is torch.dtype:
 254        v = getattr(torch, str(json_val).split(".", 1)[-1], None)
 255        return v if isinstance(v, torch.dtype) else json_val
 256    if type_hint is callable:
 257        v = _resolve_dotted_name(json_val) if isinstance(json_val, str) else None
 258        return v if (v and callable(v)) else json_val
 259    if type_hint == [type]:  # list whose elements are class objects
 260        return [
 261            (_resolve_dotted_name(v) if isinstance(v, str) else v)
 262            for v in (json_val or [])
 263        ]
 264    return json_val  # bool, int, float, str, list — JSON value is already correct
 265
 266
 267class PAIConfig:
 268    """Configuration class for PAI settings.
 269
 270    This class manages all configuration parameters for the Perforated AI system,
 271    including device settings, dendrite behavior, module conversion rules,
 272    training parameters, and debugging options.
 273
 274    Attributes
 275    ----------
 276    use_cuda : bool
 277        Whether CUDA is available and should be used.
 278    device : torch.device
 279        The device to use for computation (CPU, CUDA, etc.).
 280    save_name : str
 281        Name used for saving models (should not be set manually).
 282    debugging_output_dimensions : int
 283        Debug level for input dimension checking.
 284    confirm_correct_sizes : bool
 285        Whether to verify tensor sizes during execution.
 286    unwrapped_modules_confirmed : bool
 287        Confirmation flag for using unwrapped modules.
 288    weight_decay_accepted : bool
 289        Confirmation flag for accepting weight decay.
 290    checked_skipped_modules : bool
 291        Whether skipped modules have been verified.
 292    verbose : bool
 293        Enable verbose logging output.
 294    extra_verbose : bool
 295        Enable extra verbose logging output.
 296    silent : bool
 297        Suppress all PAI print statements.
 298    save_old_graph_scores : bool
 299        Whether to save historical graph scores.
 300    testing_dendrite_capacity : bool
 301        Enable dendrite capacity testing mode.
 302    using_safe_tensors : bool
 303        Use safe tensors file format for saving.
 304    global_candidates : int
 305        Number of global candidate dendrites.
 306    drawing_pai : bool
 307        Enable PAI visualization graphs.
 308    test_saves : bool
 309        Save intermediary test models.
 310    pai_saves : bool
 311        Save PAI-specific format models.
 312    output_dimensions : list
 313        Format specification for input tensor dimensions.
 314    improvement_threshold : float
 315        Relative improvement threshold for validation scores.
 316    improvement_threshold_raw : float
 317        Absolute improvement threshold for validation scores.
 318    candidate_weight_initialization_multiplier : float
 319        Multiplier for random dendrite weight initialization.
 320    DOING_SWITCH_EVERY_TIME : int
 321        Constant for switch mode: add dendrites every epoch.
 322    DOING_HISTORY : int
 323        Constant for switch mode: add dendrites based on validation history.
 324    n_epochs_to_switch : int
 325        Number of epochs without improvement before switching.
 326    history_lookback : int
 327        Number of epochs to average for validation history.
 328    initial_history_after_switches : int
 329        Epochs to wait after adding dendrites before beggining checks.
 330    DOING_FIXED_SWITCH : int
 331        Constant for switch mode: add dendrites at fixed intervals.
 332    fixed_switch_num : int
 333        Number of epochs between fixed switches.
 334    first_fixed_switch_num : int
 335        Number of epochs before first switch (for pretraining).
 336    DOING_NO_SWITCH : int
 337        Constant for switch mode: never add dendrites.
 338    switch_mode : int
 339        Current switch mode setting.
 340    reset_best_score_on_switch : bool
 341        Whether to reset best score when adding dendrites.
 342    learn_dendrites_live : bool
 343        Enable live dendrite learning (advanced feature).
 344    no_extra_n_modes : bool
 345        Disable extra neuron modes (advanced feature).
 346    d_type : torch.dtype
 347        Data type for dendrite weights.
 348    retain_all_dendrites : bool
 349        Keep dendrites even if they don't improve performance.
 350    find_best_lr : bool
 351        Automatically sweep learning rates when adding dendrites.
 352    dont_give_up_unless_learning_rate_lowered : bool
 353        Ensure search lowers learning rate at least once.
 354    max_dendrite_tries : int
 355        Maximum attempts to add dendrites with random initializations.
 356    max_dendrites : int
 357        Maximum total number of dendrites to add.
 358    PARAM_VALS_BY_TOTAL_EPOCH : int
 359        Constant: scheduler params tracked by total epochs.
 360    PARAM_VALS_BY_UPDATE_EPOCH : int
 361        Constant: scheduler params reset at each switch.
 362    PARAM_VALS_BY_NEURON_EPOCH_START : int
 363        Constant: scheduler params reset for neuron starts only.
 364    param_vals_setting : int
 365        Current parameter tracking mode.
 366    pai_forward_function : callable
 367        Activation function used for dendrites.
 368    modules_to_perforate : list
 369        Module types to convert to PAI modules for perforation.
 370    module_names_to_perforate : list
 371        Module names to convert to PAI modules for perforation.
 372    module_ids_to_perforate : list
 373        Specific module IDs to convert to PAI modules for perforation.
 374    modules_to_track : list
 375        Module types to track but not convert.
 376    module_names_to_track : list
 377        Module names to track but not convert.
 378    module_ids_to_track : list
 379        Specific module IDs to track but not convert.
 380    modules_to_replace : list
 381        Module types to replace before conversion.
 382    replacement_modules : list
 383        Replacement modules for modules_to_replace.
 384    modules_with_processing : list
 385        Module types requiring custom processing.
 386    modules_processing_classes : list
 387        Processing classes for modules_with_processing.
 388    module_names_with_processing : list
 389        Module names requiring custom processing.
 390    module_by_name_processing_classes : list
 391        Processing classes for module_names_with_processing.
 392    module_names_to_not_save : list
 393        Module names to exclude from saving.
 394    perforated_backpropagation : bool
 395        Whether Perforated Backpropagation is enabled.
 396    """
 397
 398    # Explicit type map for every config variable — used by load_config to coerce JSON values.
 399    # [type] means "list whose elements are class/type objects" (need dotted-name resolution).
 400    _TYPES: dict = {
 401        **{
 402            k: bool
 403            for k in (
 404                "use_cuda",
 405                "confirm_correct_sizes",
 406                "unwrapped_modules_confirmed",
 407                "weight_decay_accepted",
 408                "checked_skipped_modules",
 409                "verbose",
 410                "extra_verbose",
 411                "silent",
 412                "save_old_graph_scores",
 413                "testing_dendrite_capacity",
 414                "using_safe_tensors",
 415                "drawing_pai",
 416                "drawing_extra_graphs",
 417                "test_saves",
 418                "pai_saves",
 419                "reset_best_score_on_switch",
 420                "learn_dendrites_live",
 421                "no_extra_n_modes",
 422                "retain_all_dendrites",
 423                "find_best_lr",
 424                "dont_give_up_unless_learning_rate_lowered",
 425                "candidate_weight_init_by_main",
 426                "perforated_backpropagation",
 427                "weight_tying_experimental",
 428                "dashboard_events_enabled",
 429                "dashboard_debug",
 430            )
 431        },
 432        **{
 433            k: int
 434            for k in (
 435                "debugging_output_dimensions",
 436                "global_candidates",
 437                "n_epochs_to_switch",
 438                "history_lookback",
 439                "initial_history_after_switches",
 440                "fixed_switch_num",
 441                "first_fixed_switch_num",
 442                "switch_mode",
 443                "max_dendrite_tries",
 444                "max_dendrites",
 445                "param_vals_setting",
 446            )
 447        },
 448        **{
 449            k: float
 450            for k in (
 451                "improvement_threshold_raw",
 452                "candidate_weight_initialization_multiplier",
 453            )
 454        },
 455        **{k: str for k in ("save_name", "library_validation_score", "dashboard_url")},
 456        "device": torch.device,
 457        "d_type": torch.dtype,
 458        "pai_forward_function": callable,
 459        **{
 460            k: list
 461            for k in (
 462                "output_dimensions",
 463                "improvement_threshold",
 464                "module_names_to_perforate",
 465                "module_ids_to_perforate",
 466                "module_names_to_track",
 467                "module_ids_to_track",
 468                "parameter_ids_to_track",
 469                "module_names_with_processing",
 470                "module_names_to_not_save",
 471                "library_extra_scores",
 472                "library_extra_scores_without_graphing",
 473            )
 474        },
 475        **{
 476            k: [type]
 477            for k in (
 478                "modules_to_perforate",
 479                "modules_to_track",
 480                "modules_to_replace",
 481                "replacement_modules",
 482                "modules_with_processing",
 483                "modules_processing_classes",
 484                "module_by_name_processing_classes",
 485            )
 486        },
 487    }
 488
 489    # Subset of _TYPES: variables that can be overridden on a per-module basis
 490    # via the Studio.  These are set outside the ``if not module_name:`` block
 491    # in :py:meth:`__init__`, so they are meaningful when constructing a
 492    # module-specific PAIConfig instance.
 493    _CUSTOMIZABLE: dict = {
 494        "verbose": bool,
 495        "extra_verbose": bool,
 496        "silent": bool,
 497        "global_candidates": int,
 498        "output_dimensions": list,
 499        "candidate_weight_initialization_multiplier": float,
 500        "candidate_weight_init_by_main": bool,
 501        "retain_all_dendrites": bool,
 502        "max_dendrites": int,
 503        "pai_forward_function": callable,
 504    }
 505
 506    def __getattr__(self, name):
 507        """Handle missing attributes gracefully, especially for PB variables.
 508
 509        Parameters
 510        ----------
 511        name : str
 512            The name of the attribute being accessed.
 513
 514        Returns
 515        -------
 516        None or raises AttributeError
 517            Returns None for missing set_ methods, raises AttributeError otherwise.
 518        """
 519        if name.startswith("set_"):
 520            print(f"Variable '{name[4:]}' does not exist.  Ignoring set attempt.")
 521            return lambda x: None
 522        if name.startswith("append_"):
 523            print(
 524                f"List Variable '{name[7:]}' does not exist.  Ignoring append attempt."
 525            )
 526            return lambda x: None
 527        if name.startswith("get_") and self.__dict__.get("_module_name") is not None:
 528            # Module-specific config: check for a per-module override stored directly in
 529            # __dict__ (written by load_config when custom JSON data was found for this
 530            # module).  This covers CUSTOMIZABLE vars that are not initialised for
 531            # per-module configs (e.g. output_dimensions, which lives inside the
 532            # ``if not module_name:`` block).
 533            private_key = f"_{name[4:]}"
 534            if private_key in self.__dict__:
 535                stored = self.__dict__[private_key]
 536                return lambda: stored
 537            # Fall back to the global pc instance for vars not set on this instance.
 538            global_getter = getattr(pc, name, None)
 539            if global_getter is not None:
 540                return global_getter
 541        raise AttributeError(
 542            f"'{self.__class__.__name__}' object has no attribute '{name}'"
 543        )
 544
 545    def __init__(self, module_name=None, module_type=None):
 546        """Initialize PAIConfig with default settings.
 547
 548        module_name=None means this is the main global config.
 549        If module_name is set this is a per-module config that loads
 550        custom settings from module_settings[module_name] (by id) or
 551        module_settings[module_type] (by type) in the JSON file.
 552        """
 553        # Must be first: prevents __getattr__ from firing for _config_file
 554        # during construction (before add_pai_config_var_functions sets it).
 555        # Also disables auto-save in setters until the end of __init__.
 556        self.__dict__["_config_file"] = None
 557        # None = global config; any string = per-module config
 558        self.__dict__["_module_name"] = module_name
 559        # Short class name of the wrapped module (e.g. 'Conv2d'), used as
 560        # a fallback lookup key when the specific name has no saved settings.
 561        self.__dict__["_module_type"] = module_type
 562
 563        if not module_name:
 564            ### Global Constants
 565            # Device configuration
 566            self.use_cuda = torch.cuda.is_available()
 567            add_pai_config_var_functions(self, "use_cuda", self.use_cuda)
 568            self.device = torch.device("cuda" if self.use_cuda else "cpu")
 569            add_pai_config_var_functions(self, "device", self.device)
 570
 571            self.save_name = ""
 572            add_pai_config_var_functions(self, "save_name", self.save_name)
 573
 574            # Debug settings
 575            self.debugging_output_dimensions = 0
 576            add_pai_config_var_functions(
 577                self, "debugging_output_dimensions", self.debugging_output_dimensions
 578            )
 579            # Debugging input tensor sizes.
 580            # This will slow things down very slightly and is not necessary but can help
 581            # catch when dimensions were not filled in correctly.
 582            self.confirm_correct_sizes = False
 583            add_pai_config_var_functions(
 584                self, "confirm_correct_sizes", self.confirm_correct_sizes
 585            )
 586
 587            # Confirmation flags for non-recommended options
 588            self.unwrapped_modules_confirmed = False
 589            add_pai_config_var_functions(
 590                self, "unwrapped_modules_confirmed", self.unwrapped_modules_confirmed
 591            )
 592            self.weight_decay_accepted = False
 593            add_pai_config_var_functions(
 594                self, "weight_decay_accepted", self.weight_decay_accepted
 595            )
 596            self.checked_skipped_modules = False
 597            add_pai_config_var_functions(
 598                self, "checked_skipped_modules", self.checked_skipped_modules
 599            )
 600            # Analysis settings
 601            self.save_old_graph_scores = True
 602            add_pai_config_var_functions(
 603                self, "save_old_graph_scores", self.save_old_graph_scores
 604            )
 605            # Testing settings
 606            self.testing_dendrite_capacity = True
 607            add_pai_config_var_functions(
 608                self, "testing_dendrite_capacity", self.testing_dendrite_capacity
 609            )
 610
 611            # File format settings
 612            self.using_safe_tensors = True
 613            add_pai_config_var_functions(
 614                self, "using_safe_tensors", self.using_safe_tensors
 615            )
 616
 617            # Checkpoint loading settings
 618            # Whether to use strict=True when loading state_dict
 619            # Set to False if loading old checkpoints that are missing new fields
 620            self.strict_loading = True
 621            add_pai_config_var_functions(
 622                self, "strict_loading", self.strict_loading
 623            )
 624
 625            # Graph and visualization settings
 626            # A graph setting which can be set to false if you want to do your own
 627            # training visualizations
 628            self.drawing_pai = True
 629            add_pai_config_var_functions(self, "drawing_pai", self.drawing_pai)
 630
 631            # Drawing extra graphs beyond the standard ones.
 632            self.drawing_extra_graphs = True
 633            add_pai_config_var_functions(
 634                self, "drawing_extra_graphs", self.drawing_extra_graphs
 635            )
 636
 637            # Saving test intermediary models, good for experimentation, bad for memory
 638            self.test_saves = True
 639            add_pai_config_var_functions(self, "test_saves", self.test_saves)
 640            # To be filled in later. pai_saves will remove some extra scaffolding for
 641            # slight memory and speed improvements
 642            self.pai_saves = False
 643            add_pai_config_var_functions(self, "pai_saves", self.pai_saves)
 644            # Improvement thresholds
 645            # Percentage improvement increase needed to call a new best validation score
 646            self.improvement_threshold = [0.001, 0.0001, 0.0]
 647            add_pai_config_var_functions(
 648                self, "improvement_threshold", self.improvement_threshold
 649            )
 650
 651            # Raw increase needed
 652            self.improvement_threshold_raw = 1e-5
 653            add_pai_config_var_functions(
 654                self, "improvement_threshold_raw", self.improvement_threshold_raw
 655            )
 656            # SWITCH MODE SETTINGS
 657
 658            # Add dendrites every time to debug implementation
 659            self.DOING_SWITCH_EVERY_TIME = 0
 660
 661            # Switch when validation hasn't improved over x epochs
 662            self.DOING_HISTORY = 1
 663            # Epochs to try before deciding to load previous best and add dendrites
 664            # Be sure this is higher than scheduler patience
 665            self.n_epochs_to_switch = 10
 666            add_pai_config_var_functions(
 667                self, "n_epochs_to_switch", self.n_epochs_to_switch
 668            )
 669            # Number to average validation scores over
 670            self.history_lookback = 1
 671            add_pai_config_var_functions(
 672                self, "history_lookback", self.history_lookback
 673            )
 674            # Amount of epochs to run after adding a new set of dendrites before checking
 675            # to add more
 676            self.initial_history_after_switches = 0
 677            add_pai_config_var_functions(
 678                self,
 679                "initial_history_after_switches",
 680                self.initial_history_after_switches,
 681            )
 682
 683            # Switch after a fixed number of epochs
 684            self.DOING_FIXED_SWITCH = 2
 685            # Number of epochs to complete before switching
 686            self.fixed_switch_num = 250
 687            add_pai_config_var_functions(
 688                self, "fixed_switch_num", self.fixed_switch_num
 689            )
 690            # An additional flag if you want your first switch to occur later than all the
 691            # rest for initial pretraining.  This is a new minimum, if its lower than
 692            # the above it will be ignored.
 693            self.first_fixed_switch_num = 1
 694            add_pai_config_var_functions(
 695                self, "first_fixed_switch_num", self.first_fixed_switch_num
 696            )
 697
 698            # A setting to not add dendrites and just do regular training
 699            # Warning, this will also never trigger training_complete
 700            self.DOING_NO_SWITCH = 3
 701
 702            # Default switch mode
 703            self.switch_mode = self.DOING_HISTORY
 704            add_pai_config_var_functions(self, "switch_mode", self.switch_mode)
 705
 706            # Reset settings
 707            # Resets score on switch
 708            # This can be useful if you need many epochs to catch up to the best score
 709            # from the previous version after adding dendrites
 710            self.reset_best_score_on_switch = True
 711            add_pai_config_var_functions(
 712                self, "reset_best_score_on_switch", self.reset_best_score_on_switch
 713            )
 714
 715            # Advanced settings
 716            # Not used in open source implementation, leave as default
 717            self.learn_dendrites_live = False
 718            add_pai_config_var_functions(
 719                self, "learn_dendrites_live", self.learn_dendrites_live
 720            )
 721            self.no_extra_n_modes = True
 722            add_pai_config_var_functions(
 723                self, "no_extra_n_modes", self.no_extra_n_modes
 724            )
 725
 726            # Data type for new modules and dendrite to dendrite / dendrite to neuron
 727            # weights
 728            self.d_type = torch.float
 729            add_pai_config_var_functions(self, "d_type", self.d_type)
 730
 731            # Learning rate management
 732            # A setting to automatically sweep over previously used learning rates when
 733            # adding new dendrites
 734            # Sometimes it's best to go back to initial LR, but often its best to start
 735            # at a lower LR
 736            self.find_best_lr = True
 737            add_pai_config_var_functions(self, "find_best_lr", self.find_best_lr)
 738            # Enforces the above even if the previous epoch didn't lower the learning rate
 739            self.dont_give_up_unless_learning_rate_lowered = True
 740            add_pai_config_var_functions(
 741                self,
 742                "dont_give_up_unless_learning_rate_lowered",
 743                self.dont_give_up_unless_learning_rate_lowered,
 744            )
 745
 746            # Dendrite attempt settings
 747            # Set to 1 if you want to quit as soon as one dendrite fails
 748            # Higher values will try new random dendrite weights this many times before
 749            # accepting that more dendrites don't improve
 750            self.max_dendrite_tries = 2
 751            add_pai_config_var_functions(
 752                self, "max_dendrite_tries", self.max_dendrite_tries
 753            )
 754
 755            # Scheduler parameter settings
 756            # Have learning rate params be by total epoch
 757            self.PARAM_VALS_BY_TOTAL_EPOCH = 0
 758            # Reset the params at every switch
 759            self.PARAM_VALS_BY_UPDATE_EPOCH = 1
 760            # Reset params for dendrite starts but not for normal restarts
 761            # Not used for open source version
 762            self.PARAM_VALS_BY_NEURON_EPOCH_START = 2
 763            # Default setting
 764            self.param_vals_setting = self.PARAM_VALS_BY_UPDATE_EPOCH
 765            add_pai_config_var_functions(
 766                self, "param_vals_setting", self.param_vals_setting
 767            )
 768            # Lists for module types and names to add dendrites to
 769            # For these lists no specifier means type, name is module name
 770            # and ids is the individual modules id, eg. model.conv2
 771            self.modules_to_perforate = []
 772            add_pai_config_var_functions(
 773                self, "modules_to_perforate", self.modules_to_perforate, list_type=True
 774            )
 775            self.module_names_to_perforate = [
 776                "PAISequential",
 777                "Conv1d",
 778                "Conv2d",
 779                "Conv3d",
 780                "Linear",
 781            ]
 782            add_pai_config_var_functions(
 783                self,
 784                "module_names_to_perforate",
 785                self.module_names_to_perforate,
 786                list_type=True,
 787            )
 788            self.module_ids_to_perforate = []
 789            add_pai_config_var_functions(
 790                self,
 791                "module_ids_to_perforate",
 792                self.module_ids_to_perforate,
 793                list_type=True,
 794            )
 795
 796            # All modules should either be perforated or tracked to ensure all modules
 797            # are accounted for
 798            self.modules_to_track = []
 799            add_pai_config_var_functions(
 800                self, "modules_to_track", self.modules_to_track, list_type=True
 801            )
 802            self.module_names_to_track = []
 803            add_pai_config_var_functions(
 804                self,
 805                "module_names_to_track",
 806                self.module_names_to_track,
 807                list_type=True,
 808            )
 809            # IDs are for if you want to pass only a single module by its assigned ID rather than the module type by name
 810            self.module_ids_to_track = []
 811            add_pai_config_var_functions(
 812                self, "module_ids_to_track", self.module_ids_to_track, list_type=True
 813            )
 814
 815            # Parameter IDs to track as neuron parameters without recursive behavior
 816            # (e.g., [".my_custom_parameter"]).
 817            self.parameter_ids_to_track = []
 818            add_pai_config_var_functions(
 819                self,
 820                "parameter_ids_to_track",
 821                self.parameter_ids_to_track,
 822                list_type=True,
 823            )
 824
 825            # Replacement modules happen before the conversion,
 826            # so replaced modules will then also be run through the conversion steps
 827            # These are for modules that need to be replaced before addition of dendrites
 828            # See the resnet example in models_perforatedai
 829            self.modules_to_replace = []
 830            add_pai_config_var_functions(
 831                self, "modules_to_replace", self.modules_to_replace, list_type=True
 832            )
 833            # Modules to replace the above modules with
 834            self.replacement_modules = []
 835            add_pai_config_var_functions(
 836                self, "replacement_modules", self.replacement_modules, list_type=True
 837            )
 838
 839            # Dendrites default to modules which are one tensor input and one tensor
 840            # output in forward()
 841            # Other modules require to be labeled as modules with processing and assigned
 842            # processing classes
 843            # This can be done by module type or module name see customization.md in API
 844            # for example
 845            self.modules_with_processing = []
 846            add_pai_config_var_functions(
 847                self,
 848                "modules_with_processing",
 849                self.modules_with_processing,
 850                list_type=True,
 851            )
 852            self.modules_processing_classes = []
 853            add_pai_config_var_functions(
 854                self,
 855                "modules_processing_classes",
 856                self.modules_processing_classes,
 857                list_type=True,
 858            )
 859            self.module_names_with_processing = []
 860            add_pai_config_var_functions(
 861                self,
 862                "module_names_with_processing",
 863                self.module_names_with_processing,
 864                list_type=True,
 865            )
 866            self.module_by_name_processing_classes = []
 867            add_pai_config_var_functions(
 868                self,
 869                "module_by_name_processing_classes",
 870                self.module_by_name_processing_classes,
 871                list_type=True,
 872            )
 873
 874            # Similarly here as above. Some huggingface models have multiple pointers to
 875            # the same modules which cause problems
 876            # If you want to only save one of the multiple pointers you can set which ones
 877            # not to save here
 878            self.module_names_to_not_save = [".base_model"]
 879            add_pai_config_var_functions(
 880                self,
 881                "module_names_to_not_save",
 882                self.module_names_to_not_save,
 883                list_type=True,
 884            )
 885
 886            # Perforated Backpropagation settings
 887            self.perforated_backpropagation = False
 888            add_pai_config_var_functions(
 889                self, "perforated_backpropagation", self.perforated_backpropagation
 890            )
 891            
 892            # This is specifically a workaround for weight tying
 893            # Not to be used for a duplicate pointer that isn't actually run twice
 894            self.weight_tying_experimental = False
 895            add_pai_config_var_functions(
 896                self, "weight_tying_experimental", self.weight_tying_experimental
 897            )
 898
 899            # Dashboard event streaming settings
 900            self.dashboard_events_enabled = False
 901            add_pai_config_var_functions(
 902                self, "dashboard_events_enabled", self.dashboard_events_enabled
 903            )
 904            self.dashboard_url = "http://localhost:3002"
 905            add_pai_config_var_functions(
 906                self, "dashboard_url", self.dashboard_url
 907            )
 908            self.dashboard_debug = False
 909            add_pai_config_var_functions(
 910                self, "dashboard_debug", self.dashboard_debug
 911            )
 912
 913            # These are settings where libraries must be doing the scoring adding to
 914            # message from your main script what metric to use
 915            self.library_validation_score = ""
 916            add_pai_config_var_functions(
 917                self, "library_validation_score", self.library_validation_score
 918            )
 919            self.library_extra_scores = []
 920            add_pai_config_var_functions(
 921                self,
 922                "library_extra_scores",
 923                self.library_extra_scores,
 924                list_type=True,
 925            )
 926            self.library_extra_scores_without_graphing = []
 927            add_pai_config_var_functions(
 928                self,
 929                "library_extra_scores_without_graphing",
 930                self.library_extra_scores_without_graphing,
 931                list_type=True,
 932            )
 933
 934            # Input dimensions needs to be set every time. It is set to what format of
 935            # planes you are expecting.
 936            # Neuron index should be set to 0, variable indexes should be set to -1.
 937            # For example, if your format is [batchsize, nodes, x, y]
 938            # output_dimensions is [-1, 0, -1, -1].
 939            # if your format is, [batchsize, time index, nodes] output_dimensions is
 940            # [-1, -1, 0]
 941            self.output_dimensions = [-1, 0, -1, -1]
 942            add_pai_config_var_functions(
 943                self, "output_dimensions", self.output_dimensions, list_type=True
 944            )
 945        # Verbosity settings
 946        self.verbose = False
 947        add_pai_config_var_functions(self, "verbose", self.verbose)
 948        self.extra_verbose = False
 949        add_pai_config_var_functions(self, "extra_verbose", self.extra_verbose)
 950        # Suppress all PAI prints
 951        self.silent = False
 952        add_pai_config_var_functions(self, "silent", self.silent)
 953
 954        # In place for future implementation options of adding multiple candidate
 955        # dendrites together
 956        self.global_candidates = 1
 957        add_pai_config_var_functions(self, "global_candidates", self.global_candidates)
 958
 959        # Weight initialization settings
 960        # Multiplier when randomizing dendrite weights
 961        self.candidate_weight_initialization_multiplier = 0.01
 962        add_pai_config_var_functions(
 963            self,
 964            "candidate_weight_initialization_multiplier",
 965            self.candidate_weight_initialization_multiplier,
 966        )
 967        # Multiplier when randomizing dendrite weights
 968        self.candidate_weight_init_by_main = False
 969        add_pai_config_var_functions(
 970            self,
 971            "candidate_weight_init_by_main",
 972            self.candidate_weight_init_by_main,
 973        )
 974
 975        # Dendrite retention settings
 976        # A setting to keep dendrites even if they do not improve scores
 977        self.retain_all_dendrites = False
 978        add_pai_config_var_functions(
 979            self, "retain_all_dendrites", self.retain_all_dendrites
 980        )
 981
 982        # Max dendrites to add even if they do continue improving scores
 983        self.max_dendrites = 100
 984        add_pai_config_var_functions(self, "max_dendrites", self.max_dendrites)
 985
 986        # Activation function settings
 987        # The activation function to use for dendrites
 988        self.pai_forward_function = torch.sigmoid
 989        add_pai_config_var_functions(
 990            self, "pai_forward_function", self.pai_forward_function
 991        )
 992
 993        # ------------------------------------------------------------------
 994        # Config file will be set when save_name is assigned (in perforate_model)
 995        # ------------------------------------------------------------------
 996        # _config_file stays None until save_name is set to a non-empty value
 997
 998    # ------------------------------------------------------------------
 999
1000    def save_config(self, filename):
1001        """Save the current PAIConfig state to a JSON file.
1002
1003        Parameters
1004        ----------
1005        filename : str
1006            Destination file path (created or overwritten).
1007
1008        Notes
1009        -----
1010        Values that are not natively JSON-serialisable (torch.device,
1011        torch.dtype, nn.Module subclasses, callables) are stored as their
1012        dotted-string representations so they can be round-tripped by
1013        :py:meth:`load_config`.
1014
1015        Returns
1016        -------
1017        None
1018            This function does not return a value.
1019        """
1020        import json
1021
1022        config_dict = {}
1023
1024        # Private-storage vars added by add_pai_config_var_functions
1025        # e.g. self._module_ids_to_perforate → key 'module_ids_to_perforate'
1026        for key, val in sorted(self.__dict__.items()):
1027            if not (key.startswith("_") and not key.startswith("__")):
1028                continue
1029            # Skip internal bookkeeping keys that must not round-trip through JSON
1030            if key in ("_config_file", "_module_name", "_module_type"):
1031                continue
1032            if callable(val):  # skip bound method refs
1033                continue
1034            clean_key = key[1:]
1035            try:
1036                config_dict[clean_key] = _serialize_pai_value(val)
1037            except Exception:
1038                config_dict[clean_key] = str(val)
1039
1040        # Plain constants (DOING_*, PARAM_VALS_BY_*, etc.)
1041        for key, val in self.__dict__.items():
1042            if key.startswith("_") or callable(val):
1043                continue
1044            if key not in config_dict:
1045                try:
1046                    config_dict[key] = _serialize_pai_value(val)
1047                except Exception:
1048                    config_dict[key] = str(val)
1049
1050        # Merge short class names from modules_to_perforate into module_names_to_perforate
1051        # so the UI (and JS) only needs to check one array.
1052        type_short_names = [
1053            cls.__name__
1054            for cls in self.__dict__.get("_modules_to_perforate", [])
1055            if isinstance(cls, type)
1056        ]
1057        existing = config_dict.get("module_names_to_perforate", [])
1058        config_dict["module_names_to_perforate"] = existing + [
1059            n for n in type_short_names if n not in existing
1060        ]
1061
1062        # Preserve any per-module settings written by the Studio frontend.
1063        _existing_ms: dict = {}
1064        try:
1065            with open(filename, "r") as _f:
1066                _existing_ms = json.load(_f).get("module_settings", {})
1067        except Exception:
1068            pass
1069        config_dict["module_settings"] = _existing_ms
1070
1071        # Publish customizable field type names so the Studio frontend can
1072        # render appropriate editors without needing to import the library.
1073        config_dict["_customizable_fields"] = {
1074            k: (v.__name__ if hasattr(v, "__name__") else str(v))
1075            for k, v in PAIConfig._CUSTOMIZABLE.items()
1076        }
1077
1078        # Ensure the directory exists before saving
1079        import os
1080
1081        os.makedirs(os.path.dirname(filename), exist_ok=True)
1082
1083        with open(filename, "w") as f:
1084            json.dump(config_dict, f, indent=2)
1085        print(f"[PAI Config] Saved {len(config_dict)} variables \u2192 {filename}")
1086
1087    def load_config(self, filename, module_name=None, module_type=None):
1088        """Load PAIConfig state from a JSON file produced by :py:meth:`save_config`.
1089
1090        If *module_name* is ``None`` (default) every serialisable variable in
1091        the file is restored on this instance.
1092
1093        If *module_name* is given the lookup priority is:
1094          1. ``module_settings[module_name]`` (exact name / id match)
1095          2. ``module_settings[module_type]``  (type-level fallback)
1096          3. No-op — the defaults already set by ``__init__`` are kept.
1097
1098        Parameters
1099        ----------
1100        filename : str
1101            Path to the JSON file to read.
1102        module_name : str, optional
1103            Display name (id) of the module whose custom settings should be loaded.
1104        module_type : str, optional
1105            Short class name of the module type, used as a fallback key.
1106
1107        Returns
1108        -------
1109        None
1110            This function does not return a value.
1111        """
1112        import json
1113
1114        with open(filename, "r") as f:
1115            config_dict = json.load(f)
1116
1117        if module_name is not None:
1118            # ── Per-module load ──────────────────────────────────────────────
1119            module_settings = config_dict.get("module_settings", {})
1120            # Priority: exact name → type fallback → no-op
1121            if module_name in module_settings:
1122                custom = module_settings[module_name]
1123                resolved_key = module_name
1124            elif module_type and module_type in module_settings:
1125                custom = module_settings[module_type]
1126                resolved_key = module_type
1127            else:
1128                # No custom settings for this module or type — keep defaults.
1129                return
1130            loaded = 0
1131            skipped = 0
1132            for key, json_val in custom.items():
1133                if key not in PAIConfig._CUSTOMIZABLE:
1134                    continue
1135                type_hint = PAIConfig._TYPES.get(key)
1136                private_key = f"_{key}"
1137                # Write directly to __dict__ so we bypass any setter guards and also
1138                # correctly handle vars that are not pre-initialised on per-module
1139                # configs (e.g. output_dimensions, which only lives inside the
1140                # ``if not module_name:`` block of __init__).
1141                try:
1142                    self.__dict__[private_key] = (
1143                        _deserialize_pai_value(json_val, type_hint)
1144                        if type_hint is not None
1145                        else json_val
1146                    )
1147                    loaded += 1
1148                except Exception as exc:
1149                    print(
1150                        f"[PAI Config] Warning: could not load '{key}' for '{resolved_key}': {exc}"
1151                    )
1152                    skipped += 1
1153            print(
1154                f"[PAI Config] Loaded {loaded} custom vars for '{resolved_key}' from {filename}"
1155                + (f" ({skipped} skipped)" if skipped else "")
1156            )
1157            return
1158
1159        # ── Global load: every variable in the JSON ──────────────────────────
1160        loaded = 0
1161        skipped = 0
1162        for key, json_val in config_dict.items():
1163            # Skip internal bookkeeping and Studio-only metadata keys.
1164            # 'module_name' and 'module_type' must never overwrite the
1165            # instance's _module_name/_module_type (they are internal only).
1166            if key in (
1167                "config_file",
1168                "module_settings",
1169                "module_name",
1170                "module_type",
1171            ) or key.startswith("_"):
1172                continue
1173            type_hint = PAIConfig._TYPES.get(key)
1174            private_key = f"_{key}"
1175            if hasattr(self, private_key):
1176                try:
1177                    setattr(
1178                        self,
1179                        private_key,
1180                        (
1181                            _deserialize_pai_value(json_val, type_hint)
1182                            if type_hint is not None
1183                            else json_val
1184                        ),
1185                    )
1186                    loaded += 1
1187                except Exception as exc:
1188                    print(f"[PAI Config] Warning: could not load '{key}': {exc}")
1189                    skipped += 1
1190            elif hasattr(self, key) and not callable(getattr(self, key, None)):
1191                try:
1192                    setattr(self, key, json_val)
1193                    loaded += 1
1194                except Exception:
1195                    skipped += 1
1196
1197        print(
1198            f"[PAI Config] Loaded {loaded} variables from {filename}"
1199            + (f" ({skipped} skipped)" if skipped else "")
1200        )
1201
1202
1203class PAISequential(nn.Sequential):
1204    """Sequential module wrapper for PAI.
1205
1206    This wrapper takes an array of layers and creates a sequential container
1207    that is compatible with PAI's dendrite addition system. It should be used
1208    for normalization layers and can be used for final output layers.
1209
1210    Parameters
1211    ----------
1212    layer_array : list
1213        List of PyTorch nn.Module objects to be executed sequentially.
1214
1215    Examples
1216    --------
1217    >>> layers = [nn.Linear(2 * hidden_dim, seq_width),
1218    ...           nn.LayerNorm(seq_width)]
1219    >>> sequential_block = PAISequential(layers)
1220
1221    Notes
1222    -----
1223    This should be used for:
1224        - All normalization layers (LayerNorm, BatchNorm, etc.)
1225    This can be used for:
1226        - Final output layer and softmax combinations
1227    """
1228
1229    def __init__(self, layer_array):
1230        """Initialize PAISequential with a list of layers.
1231
1232        Parameters
1233        ----------
1234        layer_array : list
1235            List of PyTorch modules to execute in sequence.
1236        """
1237        super(PAISequential, self).__init__()
1238        self.model = nn.Sequential(*layer_array)
1239
1240    def forward(self, *args, **kwargs):
1241        """Forward pass through the sequential layers.
1242
1243        Parameters
1244        ----------
1245        *args
1246            Positional arguments passed to the first layer.
1247        **kwargs
1248            Keyword arguments passed to the layers.
1249
1250        Returns
1251        -------
1252        torch.Tensor
1253            Output from the final layer in the sequence.
1254        """
1255        return self.model(*args, **kwargs)
1256
1257
1258### Global objects and variables
1259
1260### Global Modules
1261pc = PAIConfig()
1262"""Global PAIConfig instance.
1263
1264This is the primary configuration object used throughout the PAI system.
1265Modify settings through this instance to control PAI behavior.
1266"""
1267
1268"""Pointer to the PAI Tracker.
1269
1270This will be populated with the PAI Tracker instance which handles
1271the addition of dendrites during training. Initially an empty list.
1272"""
1273pai_tracker = []
1274
1275pai_scaler = None
1276
1277# This will be set to true if perforated backpropagation is available
1278# Do not just set this to True without the library and a license, it will cause errors
1279try:
1280    import perforatedbp.globals_pbp as perforatedbp_globals
1281
1282    print("Building dendrites with Perforated Backpropagation")
1283
1284    pc.set_perforated_backpropagation(True)
1285    # This is default to True for open source version
1286    # But defaults to False for perforated backpropagation
1287    pc.set_no_extra_n_modes(False)
1288
1289    # Loop through the vars module's attributes and add them dynamically
1290    for var_name in dir(perforatedbp_globals):
1291        if not var_name.startswith("_"):
1292            add_pai_config_var_functions(
1293                pc, var_name, getattr(perforatedbp_globals, var_name)
1294            )
1295
1296    # Merge PBP type hints into PAIConfig._TYPES so load_config can correctly
1297    # round-trip all perforatedbp variables from JSON.
1298    if hasattr(perforatedbp_globals, "_TYPES"):
1299        PAIConfig._TYPES.update(perforatedbp_globals._TYPES)
1300
1301except ImportError:
1302    print("Building dendrites without Perforated Backpropagation")