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 __getstate__(self):
 507        """Tell pickle what to save when this object is serialized (e.g. torch.save).
 508
 509        The get_*/set_*/append_* methods attached to each PAIConfig instance are
 510        locally-defined closures and cannot be pickled.  Only the underlying data
 511        values (stored as _<name> in __dict__) need to be saved; __setstate__
 512        will recreate the methods by calling __init__ on load.
 513        """
 514        import types
 515
 516        # Walk every attribute on this instance and keep only the plain data values.
 517        # Bound methods (the get_*/set_*/append_* closures) are skipped because
 518        # they cannot be serialized by pickle.
 519        pickle_safe_state = {}
 520        for attr_name, attr_value in self.__dict__.items():
 521            if not isinstance(attr_value, types.MethodType):
 522                pickle_safe_state[attr_name] = attr_value
 523
 524        return pickle_safe_state
 525
 526    def __setstate__(self, saved_state):
 527        """Restore this object from a pickled state (e.g. torch.load).
 528
 529        Calls __init__ to rebuild all the get_*/set_*/append_* methods, then
 530        overlays the saved data values on top so all settings are preserved.
 531        """
 532        # Grab the identity fields needed to reconstruct the right flavor of config
 533        # (global config vs. per-module config).
 534        saved_module_name = saved_state.get("_module_name")
 535        saved_module_type = saved_state.get("_module_type")
 536
 537        # Re-run __init__ so all the get_*/set_*/append_* methods exist again.
 538        # This also sets default values for every variable, which we overwrite below.
 539        self.__init__(module_name=saved_module_name, module_type=saved_module_type)
 540
 541        # Overwrite the defaults with the actual saved values.
 542        # Only restore private data attributes (prefixed with _); skip the identity
 543        # fields we already handled above, and skip _config_file for now so that
 544        # the setter calls above do not trigger auto-save during the restore.
 545        skip_on_first_pass = {"_module_name", "_module_type", "_config_file"}
 546        for attr_name, attr_value in saved_state.items():
 547            if attr_name.startswith("_") and attr_name not in skip_on_first_pass:
 548                self.__dict__[attr_name] = attr_value
 549
 550        # Restore the config file path last.  Auto-save is driven by _config_file
 551        # being set, so we only put it back after all values are already in place.
 552        if saved_state.get("_config_file"):
 553            self.__dict__["_config_file"] = saved_state["_config_file"]
 554
 555    def __getattr__(self, name):
 556        """Handle missing attributes gracefully, especially for PB variables.
 557
 558        Parameters
 559        ----------
 560        name : str
 561            The name of the attribute being accessed.
 562
 563        Returns
 564        -------
 565        None or raises AttributeError
 566            Returns None for missing set_ methods, raises AttributeError otherwise.
 567        """
 568        if name.startswith("set_"):
 569            print(f"Variable '{name[4:]}' does not exist.  Ignoring set attempt.")
 570            return lambda x: None
 571        if name.startswith("append_"):
 572            print(
 573                f"List Variable '{name[7:]}' does not exist.  Ignoring append attempt."
 574            )
 575            return lambda x: None
 576        if name.startswith("get_") and self.__dict__.get("_module_name") is not None:
 577            # Module-specific config: check for a per-module override stored directly in
 578            # __dict__ (written by load_config when custom JSON data was found for this
 579            # module).  This covers CUSTOMIZABLE vars that are not initialised for
 580            # per-module configs (e.g. output_dimensions, which lives inside the
 581            # ``if not module_name:`` block).
 582            private_key = f"_{name[4:]}"
 583            if private_key in self.__dict__:
 584                stored = self.__dict__[private_key]
 585                return lambda: stored
 586            # Fall back to the global pc instance for vars not set on this instance.
 587            global_getter = getattr(pc, name, None)
 588            if global_getter is not None:
 589                return global_getter
 590        raise AttributeError(
 591            f"'{self.__class__.__name__}' object has no attribute '{name}'"
 592        )
 593
 594    def __init__(self, module_name=None, module_type=None):
 595        """Initialize PAIConfig with default settings.
 596
 597        module_name=None means this is the main global config.
 598        If module_name is set this is a per-module config that loads
 599        custom settings from module_settings[module_name] (by id) or
 600        module_settings[module_type] (by type) in the JSON file.
 601        """
 602        # Must be first: prevents __getattr__ from firing for _config_file
 603        # during construction (before add_pai_config_var_functions sets it).
 604        # Also disables auto-save in setters until the end of __init__.
 605        self.__dict__["_config_file"] = None
 606        # None = global config; any string = per-module config
 607        self.__dict__["_module_name"] = module_name
 608        # Short class name of the wrapped module (e.g. 'Conv2d'), used as
 609        # a fallback lookup key when the specific name has no saved settings.
 610        self.__dict__["_module_type"] = module_type
 611
 612        if not module_name:
 613            ### Global Constants
 614            # Device configuration
 615            self.use_cuda = torch.cuda.is_available()
 616            add_pai_config_var_functions(self, "use_cuda", self.use_cuda)
 617            self.device = torch.device("cuda" if self.use_cuda else "cpu")
 618            add_pai_config_var_functions(self, "device", self.device)
 619
 620            self.save_name = ""
 621            add_pai_config_var_functions(self, "save_name", self.save_name)
 622
 623            # Debug settings
 624            self.debugging_output_dimensions = 0
 625            add_pai_config_var_functions(
 626                self, "debugging_output_dimensions", self.debugging_output_dimensions
 627            )
 628            # Debugging input tensor sizes.
 629            # This will slow things down very slightly and is not necessary but can help
 630            # catch when dimensions were not filled in correctly.
 631            self.confirm_correct_sizes = False
 632            add_pai_config_var_functions(
 633                self, "confirm_correct_sizes", self.confirm_correct_sizes
 634            )
 635
 636            # Confirmation flags for non-recommended options
 637            self.unwrapped_modules_confirmed = False
 638            add_pai_config_var_functions(
 639                self, "unwrapped_modules_confirmed", self.unwrapped_modules_confirmed
 640            )
 641            self.weight_decay_accepted = False
 642            add_pai_config_var_functions(
 643                self, "weight_decay_accepted", self.weight_decay_accepted
 644            )
 645            self.checked_skipped_modules = False
 646            add_pai_config_var_functions(
 647                self, "checked_skipped_modules", self.checked_skipped_modules
 648            )
 649            # Analysis settings
 650            self.save_old_graph_scores = True
 651            add_pai_config_var_functions(
 652                self, "save_old_graph_scores", self.save_old_graph_scores
 653            )
 654            # Testing settings
 655            self.testing_dendrite_capacity = True
 656            add_pai_config_var_functions(
 657                self, "testing_dendrite_capacity", self.testing_dendrite_capacity
 658            )
 659
 660            # File format settings
 661            self.using_safe_tensors = True
 662            add_pai_config_var_functions(
 663                self, "using_safe_tensors", self.using_safe_tensors
 664            )
 665
 666            # Checkpoint loading settings
 667            # Whether to use strict=True when loading state_dict
 668            # Set to False if loading old checkpoints that are missing new fields
 669            self.strict_loading = True
 670            add_pai_config_var_functions(
 671                self, "strict_loading", self.strict_loading
 672            )
 673
 674            # Graph and visualization settings
 675            # A graph setting which can be set to false if you want to do your own
 676            # training visualizations
 677            self.drawing_pai = True
 678            add_pai_config_var_functions(self, "drawing_pai", self.drawing_pai)
 679
 680            # Drawing extra graphs beyond the standard ones.
 681            self.drawing_extra_graphs = True
 682            add_pai_config_var_functions(
 683                self, "drawing_extra_graphs", self.drawing_extra_graphs
 684            )
 685
 686            # Saving test intermediary models, good for experimentation, bad for memory
 687            self.test_saves = True
 688            add_pai_config_var_functions(self, "test_saves", self.test_saves)
 689            # To be filled in later. pai_saves will remove some extra scaffolding for
 690            # slight memory and speed improvements
 691            self.pai_saves = False
 692            add_pai_config_var_functions(self, "pai_saves", self.pai_saves)
 693            # Improvement thresholds
 694            # Percentage improvement increase needed to call a new best validation score
 695            self.improvement_threshold = [0.001, 0.0001, 0.0]
 696            add_pai_config_var_functions(
 697                self, "improvement_threshold", self.improvement_threshold
 698            )
 699
 700            # Raw increase needed
 701            self.improvement_threshold_raw = 1e-5
 702            add_pai_config_var_functions(
 703                self, "improvement_threshold_raw", self.improvement_threshold_raw
 704            )
 705            # SWITCH MODE SETTINGS
 706
 707            # Add dendrites every time to debug implementation
 708            self.DOING_SWITCH_EVERY_TIME = 0
 709
 710            # Switch when validation hasn't improved over x epochs
 711            self.DOING_HISTORY = 1
 712            # Epochs to try before deciding to load previous best and add dendrites
 713            # Be sure this is higher than scheduler patience
 714            self.n_epochs_to_switch = 10
 715            add_pai_config_var_functions(
 716                self, "n_epochs_to_switch", self.n_epochs_to_switch
 717            )
 718            # Number to average validation scores over
 719            self.history_lookback = 1
 720            add_pai_config_var_functions(
 721                self, "history_lookback", self.history_lookback
 722            )
 723            # Amount of epochs to run after adding a new set of dendrites before checking
 724            # to add more
 725            self.initial_history_after_switches = 0
 726            add_pai_config_var_functions(
 727                self,
 728                "initial_history_after_switches",
 729                self.initial_history_after_switches,
 730            )
 731
 732            # Switch after a fixed number of epochs
 733            self.DOING_FIXED_SWITCH = 2
 734            # Number of epochs to complete before switching
 735            self.fixed_switch_num = 250
 736            add_pai_config_var_functions(
 737                self, "fixed_switch_num", self.fixed_switch_num
 738            )
 739            # An additional flag if you want your first switch to occur later than all the
 740            # rest for initial pretraining.  This is a new minimum, if its lower than
 741            # the above it will be ignored.
 742            self.first_fixed_switch_num = 1
 743            add_pai_config_var_functions(
 744                self, "first_fixed_switch_num", self.first_fixed_switch_num
 745            )
 746
 747            # A setting to not add dendrites and just do regular training
 748            # Warning, this will also never trigger training_complete
 749            self.DOING_NO_SWITCH = 3
 750
 751            # Default switch mode
 752            self.switch_mode = self.DOING_HISTORY
 753            add_pai_config_var_functions(self, "switch_mode", self.switch_mode)
 754
 755            # Reset settings
 756            # Resets score on switch
 757            # This can be useful if you need many epochs to catch up to the best score
 758            # from the previous version after adding dendrites
 759            self.reset_best_score_on_switch = True
 760            add_pai_config_var_functions(
 761                self, "reset_best_score_on_switch", self.reset_best_score_on_switch
 762            )
 763
 764            # Advanced settings
 765            # Not used in open source implementation, leave as default
 766            self.learn_dendrites_live = False
 767            add_pai_config_var_functions(
 768                self, "learn_dendrites_live", self.learn_dendrites_live
 769            )
 770            self.no_extra_n_modes = True
 771            add_pai_config_var_functions(
 772                self, "no_extra_n_modes", self.no_extra_n_modes
 773            )
 774
 775            # Data type for new modules and dendrite to dendrite / dendrite to neuron
 776            # weights
 777            self.d_type = torch.float
 778            add_pai_config_var_functions(self, "d_type", self.d_type)
 779
 780            # Learning rate management
 781            # A setting to automatically sweep over previously used learning rates when
 782            # adding new dendrites
 783            # Sometimes it's best to go back to initial LR, but often its best to start
 784            # at a lower LR
 785            self.find_best_lr = True
 786            add_pai_config_var_functions(self, "find_best_lr", self.find_best_lr)
 787            # Enforces the above even if the previous epoch didn't lower the learning rate
 788            self.dont_give_up_unless_learning_rate_lowered = True
 789            add_pai_config_var_functions(
 790                self,
 791                "dont_give_up_unless_learning_rate_lowered",
 792                self.dont_give_up_unless_learning_rate_lowered,
 793            )
 794
 795            # Dendrite attempt settings
 796            # Set to 1 if you want to quit as soon as one dendrite fails
 797            # Higher values will try new random dendrite weights this many times before
 798            # accepting that more dendrites don't improve
 799            self.max_dendrite_tries = 2
 800            add_pai_config_var_functions(
 801                self, "max_dendrite_tries", self.max_dendrite_tries
 802            )
 803
 804            # Scheduler parameter settings
 805            # Have learning rate params be by total epoch
 806            self.PARAM_VALS_BY_TOTAL_EPOCH = 0
 807            # Reset the params at every switch
 808            self.PARAM_VALS_BY_UPDATE_EPOCH = 1
 809            # Reset params for dendrite starts but not for normal restarts
 810            # Not used for open source version
 811            self.PARAM_VALS_BY_NEURON_EPOCH_START = 2
 812            # Default setting
 813            self.param_vals_setting = self.PARAM_VALS_BY_UPDATE_EPOCH
 814            add_pai_config_var_functions(
 815                self, "param_vals_setting", self.param_vals_setting
 816            )
 817            # Lists for module types and names to add dendrites to
 818            # For these lists no specifier means type, name is module name
 819            # and ids is the individual modules id, eg. model.conv2
 820            self.modules_to_perforate = []
 821            add_pai_config_var_functions(
 822                self, "modules_to_perforate", self.modules_to_perforate, list_type=True
 823            )
 824            self.module_names_to_perforate = [
 825                "PAISequential",
 826                "Conv1d",
 827                "Conv2d",
 828                "Conv3d",
 829                "Linear",
 830            ]
 831            add_pai_config_var_functions(
 832                self,
 833                "module_names_to_perforate",
 834                self.module_names_to_perforate,
 835                list_type=True,
 836            )
 837            self.module_ids_to_perforate = []
 838            add_pai_config_var_functions(
 839                self,
 840                "module_ids_to_perforate",
 841                self.module_ids_to_perforate,
 842                list_type=True,
 843            )
 844
 845            # All modules should either be perforated or tracked to ensure all modules
 846            # are accounted for
 847            self.modules_to_track = []
 848            add_pai_config_var_functions(
 849                self, "modules_to_track", self.modules_to_track, list_type=True
 850            )
 851            self.module_names_to_track = []
 852            add_pai_config_var_functions(
 853                self,
 854                "module_names_to_track",
 855                self.module_names_to_track,
 856                list_type=True,
 857            )
 858            # IDs are for if you want to pass only a single module by its assigned ID rather than the module type by name
 859            self.module_ids_to_track = []
 860            add_pai_config_var_functions(
 861                self, "module_ids_to_track", self.module_ids_to_track, list_type=True
 862            )
 863
 864            # Parameter IDs to track as neuron parameters without recursive behavior
 865            # (e.g., [".my_custom_parameter"]).
 866            self.parameter_ids_to_track = []
 867            add_pai_config_var_functions(
 868                self,
 869                "parameter_ids_to_track",
 870                self.parameter_ids_to_track,
 871                list_type=True,
 872            )
 873
 874            # Replacement modules happen before the conversion,
 875            # so replaced modules will then also be run through the conversion steps
 876            # These are for modules that need to be replaced before addition of dendrites
 877            # See the resnet example in models_perforatedai
 878            self.modules_to_replace = []
 879            add_pai_config_var_functions(
 880                self, "modules_to_replace", self.modules_to_replace, list_type=True
 881            )
 882            # Modules to replace the above modules with
 883            self.replacement_modules = []
 884            add_pai_config_var_functions(
 885                self, "replacement_modules", self.replacement_modules, list_type=True
 886            )
 887
 888            # Dendrites default to modules which are one tensor input and one tensor
 889            # output in forward()
 890            # Other modules require to be labeled as modules with processing and assigned
 891            # processing classes
 892            # This can be done by module type or module name see customization.md in API
 893            # for example
 894            self.modules_with_processing = []
 895            add_pai_config_var_functions(
 896                self,
 897                "modules_with_processing",
 898                self.modules_with_processing,
 899                list_type=True,
 900            )
 901            self.modules_processing_classes = []
 902            add_pai_config_var_functions(
 903                self,
 904                "modules_processing_classes",
 905                self.modules_processing_classes,
 906                list_type=True,
 907            )
 908            self.module_names_with_processing = []
 909            add_pai_config_var_functions(
 910                self,
 911                "module_names_with_processing",
 912                self.module_names_with_processing,
 913                list_type=True,
 914            )
 915            self.module_by_name_processing_classes = []
 916            add_pai_config_var_functions(
 917                self,
 918                "module_by_name_processing_classes",
 919                self.module_by_name_processing_classes,
 920                list_type=True,
 921            )
 922
 923            # Similarly here as above. Some huggingface models have multiple pointers to
 924            # the same modules which cause problems
 925            # If you want to only save one of the multiple pointers you can set which ones
 926            # not to save here
 927            self.module_names_to_not_save = [".base_model"]
 928            add_pai_config_var_functions(
 929                self,
 930                "module_names_to_not_save",
 931                self.module_names_to_not_save,
 932                list_type=True,
 933            )
 934
 935            # Perforated Backpropagation settings
 936            self.perforated_backpropagation = False
 937            add_pai_config_var_functions(
 938                self, "perforated_backpropagation", self.perforated_backpropagation
 939            )
 940            
 941            # This is specifically a workaround for weight tying
 942            # Not to be used for a duplicate pointer that isn't actually run twice
 943            self.weight_tying_experimental = False
 944            add_pai_config_var_functions(
 945                self, "weight_tying_experimental", self.weight_tying_experimental
 946            )
 947
 948            # Dashboard event streaming settings
 949            self.dashboard_events_enabled = False
 950            add_pai_config_var_functions(
 951                self, "dashboard_events_enabled", self.dashboard_events_enabled
 952            )
 953            self.dashboard_url = "http://localhost:3002"
 954            add_pai_config_var_functions(
 955                self, "dashboard_url", self.dashboard_url
 956            )
 957            self.dashboard_debug = False
 958            add_pai_config_var_functions(
 959                self, "dashboard_debug", self.dashboard_debug
 960            )
 961
 962            # These are settings where libraries must be doing the scoring adding to
 963            # message from your main script what metric to use
 964            self.library_validation_score = ""
 965            add_pai_config_var_functions(
 966                self, "library_validation_score", self.library_validation_score
 967            )
 968            self.library_extra_scores = []
 969            add_pai_config_var_functions(
 970                self,
 971                "library_extra_scores",
 972                self.library_extra_scores,
 973                list_type=True,
 974            )
 975            self.library_extra_scores_without_graphing = []
 976            add_pai_config_var_functions(
 977                self,
 978                "library_extra_scores_without_graphing",
 979                self.library_extra_scores_without_graphing,
 980                list_type=True,
 981            )
 982
 983            # Input dimensions needs to be set every time. It is set to what format of
 984            # planes you are expecting.
 985            # Neuron index should be set to 0, variable indexes should be set to -1.
 986            # For example, if your format is [batchsize, nodes, x, y]
 987            # output_dimensions is [-1, 0, -1, -1].
 988            # if your format is, [batchsize, time index, nodes] output_dimensions is
 989            # [-1, -1, 0]
 990            self.output_dimensions = [-1, 0, -1, -1]
 991            add_pai_config_var_functions(
 992                self, "output_dimensions", self.output_dimensions, list_type=True
 993            )
 994        # Verbosity settings
 995        self.verbose = False
 996        add_pai_config_var_functions(self, "verbose", self.verbose)
 997        self.extra_verbose = False
 998        add_pai_config_var_functions(self, "extra_verbose", self.extra_verbose)
 999        # Suppress all PAI prints
1000        self.silent = False
1001        add_pai_config_var_functions(self, "silent", self.silent)
1002
1003        # In place for future implementation options of adding multiple candidate
1004        # dendrites together
1005        self.global_candidates = 1
1006        add_pai_config_var_functions(self, "global_candidates", self.global_candidates)
1007
1008        # Weight initialization settings
1009        # Multiplier when randomizing dendrite weights
1010        self.candidate_weight_initialization_multiplier = 0.01
1011        add_pai_config_var_functions(
1012            self,
1013            "candidate_weight_initialization_multiplier",
1014            self.candidate_weight_initialization_multiplier,
1015        )
1016        # Multiplier when randomizing dendrite weights
1017        self.candidate_weight_init_by_main = False
1018        add_pai_config_var_functions(
1019            self,
1020            "candidate_weight_init_by_main",
1021            self.candidate_weight_init_by_main,
1022        )
1023
1024        # Dendrite retention settings
1025        # A setting to keep dendrites even if they do not improve scores
1026        self.retain_all_dendrites = False
1027        add_pai_config_var_functions(
1028            self, "retain_all_dendrites", self.retain_all_dendrites
1029        )
1030
1031        # Max dendrites to add even if they do continue improving scores
1032        self.max_dendrites = 100
1033        add_pai_config_var_functions(self, "max_dendrites", self.max_dendrites)
1034
1035        # Activation function settings
1036        # The activation function to use for dendrites
1037        self.pai_forward_function = torch.sigmoid
1038        add_pai_config_var_functions(
1039            self, "pai_forward_function", self.pai_forward_function
1040        )
1041
1042        # ------------------------------------------------------------------
1043        # Config file will be set when save_name is assigned (in perforate_model)
1044        # ------------------------------------------------------------------
1045        # _config_file stays None until save_name is set to a non-empty value
1046
1047    # ------------------------------------------------------------------
1048
1049    def save_config(self, filename):
1050        """Save the current PAIConfig state to a JSON file.
1051
1052        Parameters
1053        ----------
1054        filename : str
1055            Destination file path (created or overwritten).
1056
1057        Notes
1058        -----
1059        Values that are not natively JSON-serialisable (torch.device,
1060        torch.dtype, nn.Module subclasses, callables) are stored as their
1061        dotted-string representations so they can be round-tripped by
1062        :py:meth:`load_config`.
1063
1064        Returns
1065        -------
1066        None
1067            This function does not return a value.
1068        """
1069        import json
1070
1071        config_dict = {}
1072
1073        # Private-storage vars added by add_pai_config_var_functions
1074        # e.g. self._module_ids_to_perforate → key 'module_ids_to_perforate'
1075        for key, val in sorted(self.__dict__.items()):
1076            if not (key.startswith("_") and not key.startswith("__")):
1077                continue
1078            # Skip internal bookkeeping keys that must not round-trip through JSON
1079            if key in ("_config_file", "_module_name", "_module_type"):
1080                continue
1081            if callable(val):  # skip bound method refs
1082                continue
1083            clean_key = key[1:]
1084            try:
1085                config_dict[clean_key] = _serialize_pai_value(val)
1086            except Exception:
1087                config_dict[clean_key] = str(val)
1088
1089        # Plain constants (DOING_*, PARAM_VALS_BY_*, etc.)
1090        for key, val in self.__dict__.items():
1091            if key.startswith("_") or callable(val):
1092                continue
1093            if key not in config_dict:
1094                try:
1095                    config_dict[key] = _serialize_pai_value(val)
1096                except Exception:
1097                    config_dict[key] = str(val)
1098
1099        # Merge short class names from modules_to_perforate into module_names_to_perforate
1100        # so the UI (and JS) only needs to check one array.
1101        type_short_names = [
1102            cls.__name__
1103            for cls in self.__dict__.get("_modules_to_perforate", [])
1104            if isinstance(cls, type)
1105        ]
1106        existing = config_dict.get("module_names_to_perforate", [])
1107        config_dict["module_names_to_perforate"] = existing + [
1108            n for n in type_short_names if n not in existing
1109        ]
1110
1111        # Preserve any per-module settings written by the Studio frontend.
1112        _existing_ms: dict = {}
1113        try:
1114            with open(filename, "r") as _f:
1115                _existing_ms = json.load(_f).get("module_settings", {})
1116        except Exception:
1117            pass
1118        config_dict["module_settings"] = _existing_ms
1119
1120        # Publish customizable field type names so the Studio frontend can
1121        # render appropriate editors without needing to import the library.
1122        config_dict["_customizable_fields"] = {
1123            k: (v.__name__ if hasattr(v, "__name__") else str(v))
1124            for k, v in PAIConfig._CUSTOMIZABLE.items()
1125        }
1126
1127        # Ensure the directory exists before saving
1128        import os
1129
1130        os.makedirs(os.path.dirname(filename), exist_ok=True)
1131
1132        with open(filename, "w") as f:
1133            json.dump(config_dict, f, indent=2)
1134        print(f"[PAI Config] Saved {len(config_dict)} variables \u2192 {filename}")
1135
1136    def load_config(self, filename, module_name=None, module_type=None):
1137        """Load PAIConfig state from a JSON file produced by :py:meth:`save_config`.
1138
1139        If *module_name* is ``None`` (default) every serialisable variable in
1140        the file is restored on this instance.
1141
1142        If *module_name* is given the lookup priority is:
1143          1. ``module_settings[module_name]`` (exact name / id match)
1144          2. ``module_settings[module_type]``  (type-level fallback)
1145          3. No-op — the defaults already set by ``__init__`` are kept.
1146
1147        Parameters
1148        ----------
1149        filename : str
1150            Path to the JSON file to read.
1151        module_name : str, optional
1152            Display name (id) of the module whose custom settings should be loaded.
1153        module_type : str, optional
1154            Short class name of the module type, used as a fallback key.
1155
1156        Returns
1157        -------
1158        None
1159            This function does not return a value.
1160        """
1161        import json
1162
1163        with open(filename, "r") as f:
1164            config_dict = json.load(f)
1165
1166        if module_name is not None:
1167            # ── Per-module load ──────────────────────────────────────────────
1168            module_settings = config_dict.get("module_settings", {})
1169            # Priority: exact name → type fallback → no-op
1170            if module_name in module_settings:
1171                custom = module_settings[module_name]
1172                resolved_key = module_name
1173            elif module_type and module_type in module_settings:
1174                custom = module_settings[module_type]
1175                resolved_key = module_type
1176            else:
1177                # No custom settings for this module or type — keep defaults.
1178                return
1179            loaded = 0
1180            skipped = 0
1181            for key, json_val in custom.items():
1182                if key not in PAIConfig._CUSTOMIZABLE:
1183                    continue
1184                type_hint = PAIConfig._TYPES.get(key)
1185                private_key = f"_{key}"
1186                # Write directly to __dict__ so we bypass any setter guards and also
1187                # correctly handle vars that are not pre-initialised on per-module
1188                # configs (e.g. output_dimensions, which only lives inside the
1189                # ``if not module_name:`` block of __init__).
1190                try:
1191                    self.__dict__[private_key] = (
1192                        _deserialize_pai_value(json_val, type_hint)
1193                        if type_hint is not None
1194                        else json_val
1195                    )
1196                    loaded += 1
1197                except Exception as exc:
1198                    print(
1199                        f"[PAI Config] Warning: could not load '{key}' for '{resolved_key}': {exc}"
1200                    )
1201                    skipped += 1
1202            print(
1203                f"[PAI Config] Loaded {loaded} custom vars for '{resolved_key}' from {filename}"
1204                + (f" ({skipped} skipped)" if skipped else "")
1205            )
1206            return
1207
1208        # ── Global load: every variable in the JSON ──────────────────────────
1209        loaded = 0
1210        skipped = 0
1211        for key, json_val in config_dict.items():
1212            # Skip internal bookkeeping and Studio-only metadata keys.
1213            # 'module_name' and 'module_type' must never overwrite the
1214            # instance's _module_name/_module_type (they are internal only).
1215            if key in (
1216                "config_file",
1217                "module_settings",
1218                "module_name",
1219                "module_type",
1220            ) or key.startswith("_"):
1221                continue
1222            type_hint = PAIConfig._TYPES.get(key)
1223            private_key = f"_{key}"
1224            if hasattr(self, private_key):
1225                try:
1226                    setattr(
1227                        self,
1228                        private_key,
1229                        (
1230                            _deserialize_pai_value(json_val, type_hint)
1231                            if type_hint is not None
1232                            else json_val
1233                        ),
1234                    )
1235                    loaded += 1
1236                except Exception as exc:
1237                    print(f"[PAI Config] Warning: could not load '{key}': {exc}")
1238                    skipped += 1
1239            elif hasattr(self, key) and not callable(getattr(self, key, None)):
1240                try:
1241                    setattr(self, key, json_val)
1242                    loaded += 1
1243                except Exception:
1244                    skipped += 1
1245
1246        print(
1247            f"[PAI Config] Loaded {loaded} variables from {filename}"
1248            + (f" ({skipped} skipped)" if skipped else "")
1249        )
1250
1251
1252class PAISequential(nn.Sequential):
1253    """Sequential module wrapper for PAI.
1254
1255    This wrapper takes an array of layers and creates a sequential container
1256    that is compatible with PAI's dendrite addition system. It should be used
1257    for normalization layers and can be used for final output layers.
1258
1259    Parameters
1260    ----------
1261    layer_array : list
1262        List of PyTorch nn.Module objects to be executed sequentially.
1263
1264    Examples
1265    --------
1266    >>> layers = [nn.Linear(2 * hidden_dim, seq_width),
1267    ...           nn.LayerNorm(seq_width)]
1268    >>> sequential_block = PAISequential(layers)
1269
1270    Notes
1271    -----
1272    This should be used for:
1273        - All normalization layers (LayerNorm, BatchNorm, etc.)
1274    This can be used for:
1275        - Final output layer and softmax combinations
1276    """
1277
1278    def __init__(self, layer_array):
1279        """Initialize PAISequential with a list of layers.
1280
1281        Parameters
1282        ----------
1283        layer_array : list
1284            List of PyTorch modules to execute in sequence.
1285        """
1286        super(PAISequential, self).__init__()
1287        self.model = nn.Sequential(*layer_array)
1288
1289    def forward(self, *args, **kwargs):
1290        """Forward pass through the sequential layers.
1291
1292        Parameters
1293        ----------
1294        *args
1295            Positional arguments passed to the first layer.
1296        **kwargs
1297            Keyword arguments passed to the layers.
1298
1299        Returns
1300        -------
1301        torch.Tensor
1302            Output from the final layer in the sequence.
1303        """
1304        return self.model(*args, **kwargs)
1305
1306
1307### Global objects and variables
1308
1309### Global Modules
1310pc = PAIConfig()
1311"""Global PAIConfig instance.
1312
1313This is the primary configuration object used throughout the PAI system.
1314Modify settings through this instance to control PAI behavior.
1315"""
1316
1317"""Pointer to the PAI Tracker.
1318
1319This will be populated with the PAI Tracker instance which handles
1320the addition of dendrites during training. Initially an empty list.
1321"""
1322pai_tracker = []
1323
1324pai_scaler = None
1325
1326# This will be set to true if perforated backpropagation is available
1327# Do not just set this to True without the library and a license, it will cause errors
1328try:
1329    import perforatedbp.globals_pbp as perforatedbp_globals
1330
1331    print("Building dendrites with Perforated Backpropagation")
1332
1333    pc.set_perforated_backpropagation(True)
1334    # This is default to True for open source version
1335    # But defaults to False for perforated backpropagation
1336    pc.set_no_extra_n_modes(False)
1337
1338    # Loop through the vars module's attributes and add them dynamically
1339    for var_name in dir(perforatedbp_globals):
1340        if not var_name.startswith("_"):
1341            add_pai_config_var_functions(
1342                pc, var_name, getattr(perforatedbp_globals, var_name)
1343            )
1344
1345    # Merge PBP type hints into PAIConfig._TYPES so load_config can correctly
1346    # round-trip all perforatedbp variables from JSON.
1347    if hasattr(perforatedbp_globals, "_TYPES"):
1348        PAIConfig._TYPES.update(perforatedbp_globals._TYPES)
1349
1350except ImportError:
1351    print("Building dendrites without Perforated Backpropagation")