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