perforatedai.modules_perforatedai

   1# Copyright (c) 2025 Perforated AI
   2
   3import copy
   4import math
   5import os
   6import pdb
   7import sys
   8import time
   9from datetime import datetime
  10
  11import numpy as np
  12import torch
  13import torch.nn as nn
  14import traceback
  15
  16from perforatedai import globals_perforatedai as GPA
  17from perforatedai import utils_perforatedai as UPA
  18
  19try:
  20    from perforatedbp import modules_pbp as MPB
  21except ModuleNotFoundError as e:
  22    # Only pass if perforatedbp package itself is missing
  23    if e.name == "perforatedbp":
  24        pass
  25    else:
  26        # perforatedbp exists but is missing a dependency
  27        raise
  28
  29
  30# Values for Dendrite training, minimally used in open source version
  31_DENDRITE_TENSOR_VALUES_BASE = [
  32    "shape"
  33]  # Shape is tensor of same shape as total neurons in module
  34_DENDRITE_SINGLE_VALUES_BASE = []
  35
  36DENDRITE_INIT_VALUES = ["initialized", "current_d_init"]
  37
  38_VALUE_TRACKER_ARRAYS_BASE = ["dendrite_outs"]
  39
  40# Cached values to avoid recomputation (each tracks its own state)
  41_cached_dendrite_tensor_values = None
  42_cached_dendrite_tensor_pb_state = None
  43_cached_dendrite_single_values = None
  44_cached_dendrite_single_pb_state = None
  45_cached_value_tracker_arrays = None
  46_cached_value_tracker_pb_state = None
  47
  48
  49def get_DENDRITE_TENSOR_VALUES():
  50    """Get DENDRITE_TENSOR_VALUES, updating from MPB if perforated_backpropagation is enabled.
  51
  52    Parameters
  53    ----------
  54    None
  55
  56    Returns
  57    -------
  58    list[str]
  59        Names of tensor attributes used for dendrite state handling.
  60    """
  61    global _cached_dendrite_tensor_values, _cached_dendrite_tensor_pb_state
  62    current_pb_state = GPA.pc.get_perforated_backpropagation()
  63
  64    if (
  65        _cached_dendrite_tensor_values is None
  66        or _cached_dendrite_tensor_pb_state != current_pb_state
  67    ):
  68        _cached_dendrite_tensor_pb_state = current_pb_state
  69        if current_pb_state:
  70            _cached_dendrite_tensor_values = MPB.update_dendrite_tensor_values(
  71                _DENDRITE_TENSOR_VALUES_BASE.copy()
  72            )
  73        else:
  74            _cached_dendrite_tensor_values = _DENDRITE_TENSOR_VALUES_BASE.copy()
  75        if current_pb_state:
  76            _cached_dendrite_tensor_values = _cached_dendrite_tensor_values + MPB._variant_tensor_values
  77
  78    return _cached_dendrite_tensor_values
  79
  80
  81def get_DENDRITE_SINGLE_VALUES():
  82    """Get DENDRITE_SINGLE_VALUES, updating from MPB if perforated_backpropagation is enabled.
  83
  84    Parameters
  85    ----------
  86    None
  87
  88    Returns
  89    -------
  90    list[str]
  91        Names of scalar attributes used for dendrite state handling.
  92    """
  93    global _cached_dendrite_single_values, _cached_dendrite_single_pb_state
  94    current_pb_state = GPA.pc.get_perforated_backpropagation()
  95
  96    if (
  97        _cached_dendrite_single_values is None
  98        or _cached_dendrite_single_pb_state != current_pb_state
  99    ):
 100        _cached_dendrite_single_pb_state = current_pb_state
 101        if current_pb_state:
 102            _cached_dendrite_single_values = MPB.update_dendrite_single_values(
 103                _DENDRITE_SINGLE_VALUES_BASE.copy()
 104            )
 105        else:
 106            _cached_dendrite_single_values = _DENDRITE_SINGLE_VALUES_BASE.copy()
 107        if current_pb_state:
 108            _cached_dendrite_single_values = _cached_dendrite_single_values + MPB._variant_single_values
 109
 110    return _cached_dendrite_single_values
 111
 112
 113def get_VALUE_TRACKER_ARRAYS():
 114    """Get VALUE_TRACKER_ARRAYS, updating from MPB if perforated_backpropagation is enabled.
 115
 116    Parameters
 117    ----------
 118    None
 119
 120    Returns
 121    -------
 122    list[str]
 123        Names of array-valued fields tracked in ``DendriteValueTracker``.
 124    """
 125    global _cached_value_tracker_arrays, _cached_value_tracker_pb_state
 126    current_pb_state = GPA.pc.get_perforated_backpropagation()
 127
 128    if (
 129        _cached_value_tracker_arrays is None
 130        or _cached_value_tracker_pb_state != current_pb_state
 131    ):
 132        _cached_value_tracker_pb_state = current_pb_state
 133        if current_pb_state:
 134            _cached_value_tracker_arrays = MPB.update_value_tracker_arrays(
 135                _VALUE_TRACKER_ARRAYS_BASE.copy()
 136            )
 137        else:
 138            _cached_value_tracker_arrays = _VALUE_TRACKER_ARRAYS_BASE.copy()
 139
 140    return _cached_value_tracker_arrays
 141
 142
 143def get_DENDRITE_REINIT_VALUES():
 144    """Get DENDRITE_REINIT_VALUES.
 145
 146    Parameters
 147    ----------
 148    None
 149
 150    Returns
 151    -------
 152    list[str]
 153        Combined list of attribute names that must be reinitialized.
 154    """
 155    return get_DENDRITE_TENSOR_VALUES() + get_DENDRITE_SINGLE_VALUES()
 156
 157
 158def get_DENDRITE_SAVE_VALUES():
 159    """Get DENDRITE_SAVE_VALUES.
 160
 161    Parameters
 162    ----------
 163    None
 164
 165    Returns
 166    -------
 167    list[str]
 168        Combined list of attribute names persisted for save/load.
 169    """
 170    return (
 171        get_DENDRITE_TENSOR_VALUES()
 172        + get_DENDRITE_SINGLE_VALUES()
 173        + DENDRITE_INIT_VALUES
 174    )
 175
 176
 177def filter_backward(grad_out, values, module=None):
 178    """Filter backward pass for gradient processing.
 179
 180    This function processes gradients during the backward pass,
 181    ensuring correct input dimensions,and applying perforated backpropagation if enabled.
 182
 183    Parameters
 184    ----------
 185    grad_out : torch.Tensor
 186        The gradient output tensor from the backward pass.
 187    values : DendriteValueTracker
 188        A DendriteValueTracker instance containing values associated with the module being processed.
 189    module : PAINeuronModule, optional
 190        The owning PAINeuronModule instance. When provided and PBP is disabled,
 191        the hook deregisters itself after the one-time initialization completes,
 192        eliminating per-batch Python overhead on all subsequent backward passes.
 193
 194    Returns
 195    -------
 196    None
 197    """
 198    if GPA.pc.get_extra_verbose():
 199        print(f"{values[0].layer_name} calling backward")
 200
 201    with torch.no_grad():
 202        val = grad_out.detach()
 203        # Fast path: Python bool set after first backward in this process.
 204        # Fallback: current_d_init is a saved buffer — survives checkpoint reload
 205        # even when _fb_init_done resets to False in a fresh process.
 206        already_init = (module is not None and module._fb_init_done) or \
 207                       values[0].current_d_init.item()
 208        if not already_init:
 209            # If input dimensions and gradient don't have same shape trigger error and quit
 210            if len(values[0].this_output_dimensions) != len(grad_out.shape):
 211                print(
 212                    "The following module has not properly set this_output_dimensions"
 213                )
 214                print(values[0].layer_name)
 215                print("it is expecting:")
 216                print(values[0].this_output_dimensions)
 217                print("but received")
 218                print(grad_out.shape)
 219                print(
 220                    "to check these all at once set GPA.pc.set_debugging_output_dimensions(1)"
 221                )
 222                print(
 223                    f"Call MODEL_VARIABLE{values[0].layer_name}.set_this_output_dimensions([...]) on this module after perforate_model"
 224                )
 225                print(
 226                    "where the ... is replaced with the correct vector as described in section 4 of customization.md"
 227                )
 228                if not GPA.pc.get_debugging_output_dimensions():
 229                    sys.exit(0)
 230                else:
 231                    GPA.pc.set_debugging_output_dimensions(2)
 232                    return
 233            # Make sure that the input dimensions are correct
 234            for i in range(len(values[0].this_output_dimensions)):
 235                if values[0].this_output_dimensions[i] == 0:
 236                    continue
 237                # Make sure all input dimensions are either -1 (reduce), 1 (retain), or exact values (old format)
 238                if (
 239                    not (grad_out.shape[i] == values[0].this_output_dimensions[i])
 240                    and not values[0].this_output_dimensions[i] == -1
 241                    and not values[0].this_output_dimensions[i] == 1
 242                ):
 243                    print(
 244                        "The following module has not properly set this_output_dimensions with this incorrect shape"
 245                    )
 246                    print(values[0].layer_name)
 247                    print("it is expecting:")
 248                    print(values[0].this_output_dimensions)
 249                    print("but received")
 250                    print(grad_out.shape)
 251                    print(
 252                        "to check these all at once set GPA.pc.set_debugging_output_dimensions(1)"
 253                    )
 254                    if not GPA.pc.get_debugging_output_dimensions():
 255                        sys.exit(0)
 256                    else:
 257                        GPA.pc.set_debugging_output_dimensions(2)
 258                        return
 259            # Setup the arrays with the now known shape
 260            with torch.no_grad():
 261                if GPA.pc.get_verbose():
 262                    print("setting d shape for")
 263                    print(values[0].layer_name)
 264                    print(val.size())
 265
 266                values[0].set_out_channels(val.size())
 267                ndim = len(values[0].this_output_dimensions)
 268                storage_shape = [1] * ndim
 269                for _i in range(ndim):
 270                    if values[0].this_output_dimensions[_i] == 1:
 271                        storage_shape[_i] = val.shape[_i]
 272                storage_shape[values[0].this_node_index.item()] = values[0].out_channels
 273                values[0].setup_arrays(storage_shape)
 274            # Flag that it has been setup (both the GPU tensor and the fast Python bool)
 275            values[0].current_d_init[0] = 1
 276            # If fixed_input_sizes is enabled, populate the tuple caches now
 277            # that val.shape is known. get_tuples_and_mult will read these on
 278            # every subsequent call instead of recomputing.
 279            if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_fixed_input_sizes():
 280                from perforatedbp import modules_pbp as _MPB
 281                math_tuple, view_tuple, full_mult = _MPB.get_tuples_and_mult(val, values[0])
 282                ndim = len(val.shape)
 283                # math_tuple can be shorter than ndim (excludes this_node_index and
 284                # retained dims). Pad with -1 sentinel to fill the ndim-length buffer.
 285                padded_math = math_tuple + [-1] * (ndim - len(math_tuple))
 286                values[0].math_tuple_cache.copy_(
 287                    torch.tensor(padded_math, dtype=torch.long, device=val.device)
 288                )
 289                values[0].view_tuple_cache.copy_(
 290                    torch.tensor(view_tuple, dtype=torch.long, device=val.device)
 291                )
 292                values[0].full_mult_cache[0] = full_mult
 293            if module is not None:
 294                module._fb_init_done = True
 295                # When PBP is disabled this hook has no further work to do.
 296                # Deregister it so it never fires again, eliminating per-batch
 297                # Python overhead on all subsequent backward passes.
 298                if not GPA.pc.get_perforated_backpropagation() and module._fb_hook_handle is not None:
 299                    module._fb_hook_handle.remove()
 300                    module._fb_hook_handle = None
 301        if GPA.pc.get_perforated_backpropagation():
 302            MPB.filter_backward_pb(val, values)
 303
 304
 305def set_wrapped_params(model):
 306    """Set parameters as wrapped with dendrites.
 307
 308    Parameters
 309    ----------
 310    model : torch.nn.Module
 311        The model whose parameters are to be marked as wrapped.
 312
 313    Returns
 314    -------
 315    None
 316
 317    """
 318    for param in model.parameters():
 319        param.wrapped = True
 320
 321
 322def set_tracked_params(model):
 323    """Set parameters as tracked without dendrites.
 324
 325    Parameters
 326    ----------
 327    model : torch.nn.Module
 328        The model whose parameters are to be marked as tracked.
 329
 330    Returns
 331    -------
 332    None
 333    """
 334    for param in model.parameters():
 335        param.tracked = True
 336
 337
 338class PAINeuronModule(nn.Module):
 339    """Wrapper to set a module as one that will have dendritic copies."""
 340
 341    def __init__(self, start_module, name):
 342        """Initialize PAINeuronModule.
 343
 344        This function sets up the neuron module to wrap the start_module
 345        and manage its dendritic connections.
 346
 347        Parameters
 348        ----------
 349        start_module : nn.Module
 350            The module to wrap.
 351        name : str
 352            The name of the neuron module.
 353        """
 354        super(PAINeuronModule, self).__init__()
 355
 356        if isinstance(start_module, nn.Module):
 357            self.main_module = start_module
 358        else:
 359            print("start_module must be nn.Module: %s" % name)
 360            print(type(start_module))
 361            print(start_module)
 362            sys.exit(-1)
 363        self.name = name
 364        # Per-module config: loads custom settings from {save_name}_config.json if present.
 365        # Passes both the instance name (id) and the module type so load_config can
 366        # fall back to type-level settings when no name-specific entry exists.
 367        _module_type_name = type(start_module).__name__
 368        self.module_config = GPA.PAIConfig(
 369            module_name=self.name, module_type=_module_type_name
 370        )
 371
 372        set_wrapped_params(self.main_module)
 373        if self.module_config.get_verbose():
 374            print(
 375                f"initing a module {self.name} with main type {type(self.main_module)}"
 376            )
 377            print(start_module)
 378
 379        # If this main_module is one that requires processing set the processor
 380        if type(self.main_module) in self.module_config.get_modules_with_processing():
 381            module_index = self.module_config.get_modules_with_processing().index(
 382                type(self.main_module)
 383            )
 384            self.processor = self.module_config.get_modules_processing_classes()[
 385                module_index
 386            ]()
 387            if self.module_config.get_verbose():
 388                print("with processor")
 389                print(self.processor)
 390        elif (
 391            type(self.main_module).__name__
 392            in self.module_config.get_module_names_with_processing()
 393        ):
 394            module_index = self.module_config.get_module_names_with_processing().index(
 395                type(self.main_module).__name__
 396            )
 397            self.processor = self.module_config.get_module_by_name_processing_classes()[
 398                module_index
 399            ]()
 400            if self.module_config.get_verbose():
 401                print("with processor")
 402                print(self.processor)
 403        else:
 404            self.processor = None
 405
 406        # Field that can be filled in if your activation function requires a parameter
 407        self.activation_function_value = -1
 408        self.type = "neuron_module"
 409
 410        self.register_buffer(
 411            "this_output_dimensions",
 412            (torch.tensor(self.module_config.get_output_dimensions())),
 413        )
 414        if (self.this_output_dimensions == 0).sum() != 1:
 415            print(f"5 Need exactly one 0 in the input dimensions: {self.name}")
 416            print(self.this_output_dimensions)
 417            sys.exit(-1)
 418        self.register_buffer(
 419            "this_node_index",
 420            torch.tensor(self.module_config.get_output_dimensions().index(0)),
 421        )
 422        self.dendrite_modules_added = 0
 423
 424        # Values for dendrite to neuron weights
 425        self.dendrites_to_top = nn.ParameterList()
 426        self.register_parameter("newest_dendrite_to_top", None)
 427        self.candidate_to_top = nn.ParameterList()
 428        self.register_parameter("current_candidate_to_top", None)
 429        # Create the dendrite module
 430        self.dendrite_module = PAIDendriteModule(
 431            self.main_module,
 432            activation_function_value=self.activation_function_value,
 433            name=self.name,
 434            output_dimensions=self.this_output_dimensions,
 435        )
 436        # If it is linear and default has convolutional dimensions, automatically set to just be batch size and neuron indexes
 437        if (
 438            issubclass(type(start_module), nn.Linear)
 439            or (
 440                issubclass(type(start_module), GPA.PAISequential)
 441                and issubclass(type(start_module.model[0]), nn.Linear)
 442            )
 443        ) and (
 444            np.array(self.this_output_dimensions)[2:] == -1
 445        ).all():  # Everything past 2 is a negative 1
 446            self.set_this_output_dimensions(self.this_output_dimensions[0:2])
 447        if (
 448            issubclass(type(start_module), nn.Conv1d)
 449            or (
 450                issubclass(type(start_module), GPA.PAISequential)
 451                and issubclass(type(start_module.model[0]), nn.Conv1d)
 452            )
 453        ) and (
 454            np.array(self.this_output_dimensions)[3:] == -1
 455        ).all():  # Everything past 2 is a negative 1
 456            self.set_this_output_dimensions(self.this_output_dimensions[0:3])
 457        # Apply per-module output_dimensions override from config if present
 458        _custom_dims = self.module_config.__dict__.get("_output_dimensions")
 459        if _custom_dims is not None:
 460            self.set_this_output_dimensions(torch.tensor(_custom_dims))
 461        GPA.pai_tracker.add_pai_neuron_module(self)
 462        if self.module_config.get_perforated_backpropagation():
 463            MPB.set_neuron_parameters(self.main_module)
 464
 465        # Track filter_backward initialization state as a
 466        # plain Python bool (avoids a GPU .item() sync on every backward pass).
 467        # _fb_hook_handle stores the registered hook handle so it can be
 468        # removed after initialization completes when PBP is disabled.
 469        self._fb_init_done = False
 470        self._fb_hook_handle = None
 471
 472    def __getattr__(self, name):
 473        """Get member variables from the main module.
 474
 475        Parameters
 476        ----------
 477        name : str
 478            The name of the variable to retrieve.
 479        Returns
 480        -------
 481        The requested variable.
 482
 483        Notes
 484        -----
 485        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
 486        If it fails, it tries to get the attribute from the wrapped main_module.
 487        This allows seamless access to the main module's attributes without modifying original code.
 488        """
 489        try:
 490            return super().__getattr__(name)
 491        except AttributeError:
 492            return getattr(self.main_module, name)
 493
 494    def __getitem__(self, index):
 495        """Support indexing operations on the main module.
 496
 497        Parameters
 498        ----------
 499        index : int or slice
 500            The index or slice to retrieve.
 501
 502        Returns
 503        -------
 504        The indexed item from the main module.
 505        """
 506        return self.main_module[index]
 507
 508    def apply_pb_grads(self):
 509        """Apply perforated backpropagation gradients if enabled.
 510
 511        Parameters
 512        ----------
 513        None
 514
 515        Returns
 516        -------
 517        None
 518            This function does not return a value.
 519        """
 520        self.dendrite_module.apply_pb_grads()
 521
 522    def apply_pb_zero(self):
 523        """Clear leftover saved tensors if there are any.
 524
 525        Parameters
 526        ----------
 527        None
 528
 529        Returns
 530        -------
 531        None
 532            This function does not return a value.
 533        """
 534        self.dendrite_module.apply_pb_zero()
 535
 536    def clear_processors(self):
 537        """Clear processors if they save values for DeepCopy and save.
 538
 539        Parameters
 540        ----------
 541        None
 542
 543        Returns
 544        -------
 545        None
 546        """
 547
 548        if not self.processor:
 549            return
 550        else:
 551            self.processor.clear_processor()
 552            self.dendrite_module.clear_processors()
 553
 554    def clear_dendrites(self):
 555        """Clear and reset dendrites before loading from a state dict.
 556
 557        Parameters
 558        ----------
 559        None
 560
 561        Returns
 562        -------
 563        None
 564
 565        """
 566        # Loading a saved state reconstructs PAIDendriteModule before simulating
 567        # its saved cycles. Preserve a registered variant factory so candidate
 568        # creation does not silently fall back to deep-copying the parent module.
 569        create_dendrite_fn = self.dendrite_module._create_dendrite_fn
 570        self.dendrite_modules_added = 0
 571        self.dendrites_to_top = nn.ParameterList()
 572        self.candidate_to_top = nn.ParameterList()
 573        self.dendrite_module = PAIDendriteModule(
 574            self.main_module,
 575            activation_function_value=self.activation_function_value,
 576            name=self.name,
 577            output_dimensions=self.this_output_dimensions,
 578        )
 579        if create_dendrite_fn is not None:
 580            self.dendrite_module.set_create_dendrite(create_dendrite_fn)
 581
 582    def __str__(self):
 583        """String representation of the module.
 584
 585        Parameters
 586        ----------
 587        None
 588
 589        Returns
 590        -------
 591        str
 592            String representation of the module.
 593
 594        Notes
 595        -----
 596        Setting for verbose changes level of details in the string output.
 597        """
 598        # If verbose print the whole module otherwise just print the module type as a PAIModule
 599        if self.module_config.get_verbose():
 600            total_string = self.main_module.__str__()
 601            total_string = "PAIModule(" + total_string + ")"
 602            return total_string + self.dendrite_module.__str__()
 603        else:
 604            total_string = self.main_module.__str__()
 605            total_string = "PAIModule(" + total_string + ")"
 606            return total_string
 607
 608    def __repr__(self):
 609        """Representation of the module."""
 610        return self.__str__()
 611
 612    def set_this_output_dimensions(self, new_output_dimensions):
 613        """Set the input dimensions for the neuron and dendrite blocks.
 614
 615        Signals to this NeuronModule that its input dimensions are different
 616        than the global default.
 617
 618        Parameters
 619        ----------
 620        new_output_dimensions : list
 621            A list or tensor specifying the new input dimensions.
 622        Returns
 623        -------
 624        None
 625
 626        """
 627        if type(new_output_dimensions) is list:
 628            new_output_dimensions = torch.tensor(new_output_dimensions)
 629        delattr(self, "this_output_dimensions")
 630        self.register_buffer(
 631            "this_output_dimensions", new_output_dimensions.detach().clone()
 632        )
 633        if (new_output_dimensions == 0).sum() != 1:
 634            print(f"6 need exactly one 0 in the input dimensions: {self.name}")
 635            print(new_output_dimensions)
 636        self.this_node_index.copy_(
 637            (new_output_dimensions == 0).nonzero(as_tuple=True)[0][0]
 638        )
 639        self.dendrite_module.set_this_output_dimensions(new_output_dimensions)
 640
 641    def set_create_dendrite(self, fn):
 642        """Set a custom function for creating dendrite modules.
 643
 644        Override how dendrite copies are created from the parent module.  By default
 645        the parent module is deep-copied.  Pass any callable with the signature
 646        ``fn(parent_module) -> nn.Module`` to replace that behaviour.
 647
 648        Parameters
 649        ----------
 650        fn : callable
 651            A function with signature ``fn(parent_module) -> nn.Module``.
 652
 653        Returns
 654        -------
 655        None
 656        """
 657        self.dendrite_module.set_create_dendrite(fn)
 658
 659
 660    def set_mode(self, mode):
 661        """Switch between neuron training and dendrite training.
 662
 663        Parameters
 664        ----------
 665        mode : str
 666            The mode to set. Either "n" for neuron training or "p" for pai-dendrite training.
 667
 668        Returns
 669        -------
 670        bool
 671            True if mode was set successfully, False otherwise.
 672
 673        Notes
 674        -----
 675        If False is returned, the mode was not changed due to an error.
 676        This is a problem that should not be ignored, but it can be ignored
 677        by calling PGA.pc.set_checked_skipped_modules(True)
 678        """
 679
 680        if self.module_config.get_verbose():
 681            print(f"{self.name} calling set mode {mode}")
 682        # If returning to neuron training
 683        if mode == "n":
 684            self.dendrite_module.set_mode(mode)
 685            # Initialize the dendrite to neuron connections
 686            if self.dendrite_modules_added > 0:
 687                if self.module_config.get_learn_dendrites_live():
 688                    values = torch.cat(
 689                        (
 690                            self.dendrites_to_top[self.dendrite_modules_added - 1],
 691                            nn.Parameter(
 692                                self.candidate_to_top.detach()
 693                                .clone()
 694                                .to(dtype=self.module_config.get_d_type())
 695                            ),
 696                        ),
 697                        0,
 698                    )
 699                else:
 700                    values = torch.cat(
 701                        (
 702                            self.dendrites_to_top[self.dendrite_modules_added - 1],
 703                            nn.Parameter(
 704                                torch.zeros(
 705                                    (1, self.out_channels),
 706                                    device=self.dendrites_to_top[
 707                                        self.dendrite_modules_added - 1
 708                                    ].device,
 709                                    dtype=self.module_config.get_d_type(),
 710                                )
 711                            ),
 712                        ),
 713                        0,
 714                    )
 715                # Freeze the previous dendrites_to_top entry
 716                # so the optimizer no longer updates stale entries. Only the
 717                # most recently appended entry is ever used in the forward pass.
 718                if len(self.dendrites_to_top) > 0:
 719                    self.dendrites_to_top[-1].requires_grad_(False)
 720                self.dendrites_to_top.append(
 721                    nn.Parameter(
 722                        values.detach()
 723                        .clone()
 724                        .to(
 725                            device=self.module_config.get_device(),
 726                            dtype=self.module_config.get_d_type(),
 727                        ),
 728                        requires_grad=True,
 729                    )
 730                )
 731            else:
 732                if self.module_config.get_learn_dendrites_live():
 733                    # Freeze previous entry before appending new one.
 734                    if len(self.dendrites_to_top) > 0:
 735                        self.dendrites_to_top[-1].requires_grad_(False)
 736                    self.dendrites_to_top.append(
 737                        nn.Parameter(
 738                            self.candidate_to_top.detach()
 739                            .clone()
 740                            .to(dtype=self.module_config.get_d_type()),
 741                            requires_grad=True,
 742                        )
 743                    )
 744                else:
 745                    # Freeze previous entry before appending new one.
 746                    if len(self.dendrites_to_top) > 0:
 747                        self.dendrites_to_top[-1].requires_grad_(False)
 748                    self.dendrites_to_top.append(
 749                        nn.Parameter(
 750                            torch.zeros(
 751                                (1, self.out_channels),
 752                                device=self.module_config.get_device(),
 753                                dtype=self.module_config.get_d_type(),
 754                            )
 755                            .detach()
 756                            .clone(),
 757                            requires_grad=True,
 758                        )
 759                    )
 760            self.dendrite_modules_added += 1
 761            if self.module_config.get_perforated_backpropagation():
 762                MPB.set_module_n_pb(self)
 763                MPB.set_neuron_parameters(self.dendrites_to_top)
 764
 765        # If starting dendrite training
 766        else:
 767            try:
 768                # Save the values that were calculated in filter_backward
 769                self.out_channels = self.dendrite_module.dendrite_values[0].out_channels
 770                self.dendrite_module.out_channels = (
 771                    self.dendrite_module.dendrite_values[0].out_channels
 772                )
 773            except Exception as e:
 774                print(e)
 775                print(
 776                    f"this occurred in module: {self.dendrite_module.dendrite_values[0].layer_name}"
 777                )
 778                print(
 779                    "Module should be added to module_names_to_track so it doesn't have dendrites added"
 780                )
 781                print("If you are getting here but out_channels has not been set")
 782                print(
 783                    "A common reason is that this module never had gradients flow through it."
 784                )
 785                print("I have seen this happen because:")
 786                print("-The weights were frozen (requires_grad = False)")
 787                print(
 788                    "-A model is added but not used so it was converted to a perforated module initialized"
 789                )
 790                print(
 791                    "-A module was converted that doesn't have weights that get modified so backward doesn't flow through it"
 792                )
 793                print(
 794                    "If this is normal behavior set GPA.pc.set_checked_skipped_modules(True) in the main to ignore"
 795                )
 796                print(
 797                    "You can also set right now in this pdb terminal to have this not happen more after checking all modules this cycle."
 798                )
 799                if not self.module_config.get_checked_skipped_modules():
 800                    pdb.set_trace()
 801                return False
 802            # Only change mode if it makes it past the above exception
 803            self.dendrite_module.set_mode(mode)
 804            if self.module_config.get_perforated_backpropagation():
 805                MPB.set_module_p_pb(self)
 806        return True
 807
 808    def create_new_dendrite_module(self):
 809        """Add an additional dendrite module.
 810
 811        Parameters
 812        ----------
 813        None
 814
 815        Returns
 816        -------
 817        None
 818        """
 819        self.dendrite_module.create_new_dendrite_module(self.main_module)
 820
 821    def forward(self, *args, **kwargs):
 822        """Forward pass through the neuron module.
 823
 824        Parameters
 825        ----------
 826        *args : tuple
 827            Positional arguments for the forward pass.
 828        **kwargs : dict
 829            Keyword arguments for the forward pass.
 830
 831        Returns
 832        -------
 833        Any
 834            The output of the module after processing through the neuron and dendrite modules.
 835
 836        Notes
 837        -----
 838            The output of this forward function will have the same format as the output
 839            of the original module
 840        """
 841
 842        # If debugging all input dimensions, quit program on first forward call
 843        if self.module_config.get_debugging_output_dimensions() == 2:
 844            print("all input dim problems now printed")
 845            sys.exit(0)
 846        if self.module_config.get_extra_verbose():
 847            print(f"{self.name} calling forward")
 848        # Call the main modules forward
 849        out = self.main_module(*args, **kwargs)
 850        # Filter with the processor if required
 851        if self.processor is not None:
 852            try:
 853                out = self.processor.post_n1(out)
 854            except Exception as e:
 855                traceback.print_exc(limit=None, chain=True)
 856                print(f"Your post_n1 processor for {self.name} caused this error")
 857                print(
 858                    f"You must check how this is defined and ensure that it is properly"
 859                )
 860                print(f"accepting outputs from the neuron module and returning the")
 861                print(f"single tensor to be combined with the dendrites output tensor")
 862                sys.exit()
 863        # Call the forwards for all of the Dendrites
 864        (
 865            dendrite_outs,
 866            candidate_outs,
 867            candidate_nonlinear_outs,
 868            candidate_outs_non_zeroed,
 869        ) = self.dendrite_module(*args, **kwargs)
 870        # If there are dendrites add all of their outputs to the neurons output
 871        if self.dendrite_modules_added > 0:
 872            for i in range(0, self.dendrite_modules_added):
 873                to_top = self.dendrites_to_top[self.dendrite_modules_added - 1][i, :]
 874                for dim in range(len(dendrite_outs[i].shape)):
 875                    if dim == self.this_node_index:
 876                        continue
 877                    to_top = to_top.unsqueeze(dim)
 878                if self.module_config.get_confirm_correct_sizes():
 879                    to_top = to_top.expand(
 880                        list(dendrite_outs[i].size())[0 : self.this_node_index]
 881                        + [self.out_channels]
 882                        + list(dendrite_outs[i].size())[self.this_node_index + 1 :]
 883                    )
 884                out = out + (dendrite_outs[i].to(out.device) * to_top.to(out.device))
 885
 886        # If learning live, add the candidate's output to the neuron's output via the live weight
 887        if self.module_config.get_perforated_backpropagation():
 888            out = MPB.apply_live_candidate_to_output(
 889                self, out, candidate_nonlinear_outs
 890            )
 891
 892        # Catch if processors are required
 893        if type(out) is tuple:
 894            print(self)
 895            print(
 896                f"The output of the above module {self.name} is a tuple when it must be a single tensor"
 897            )
 898            print(
 899                "This must be fixed to enable the dendrite and neuron output to be combined"
 900            )
 901            print(
 902                "Look in the API customization.md at section 2.2 regarding processors to fix this."
 903            )
 904            pdb.set_trace()
 905
 906        # Call filter backward to ensure the neuron index is setup correctly.
 907        # Register only when not yet initialized (PBP=False) or always when PBP
 908        # is enabled. Store the handle so the hook can deregister itself after
 909        # the one-time initialization completes (PBP=False path).
 910        if out.requires_grad and (not self._fb_init_done or GPA.pc.get_perforated_backpropagation()):
 911            self._fb_hook_handle = out.register_hook(
 912                lambda grad: filter_backward(grad, self.dendrite_module.dendrite_values, self)
 913            )
 914
 915        # If there is a processor apply the second neuron stage
 916        if self.processor is not None:
 917            try:
 918                out = self.processor.post_n2(out)
 919            except Exception as e:
 920                traceback.print_exc(limit=None, chain=True)
 921                print(f"Your post_n2 processor for {self.name} caused this error")
 922                print(
 923                    f"You must check how this is defined and ensure that it is properly"
 924                )
 925                print(
 926                    f"accepting the output tensor after combining the neuron's output "
 927                )
 928                print(f"with the dendrite's output and returning something that is the")
 929                print(f"same format as your original module's return")
 930                sys.exit()
 931        return out
 932
 933
 934class TrackedNeuronModule(nn.Module):
 935    """Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for."""
 936
 937    def __init__(self, start_module, name):
 938        """Initialize TrackedNeuronModule.
 939
 940        This function sets up the tracked neuron module to wrap the start_module
 941        without adding dendrites.
 942
 943        Parameters
 944        ----------
 945        start_module : nn.Module
 946            The module to wrap.
 947        name : str
 948            The name of the neuron module.
 949        """
 950        super(TrackedNeuronModule, self).__init__()
 951
 952        if isinstance(start_module, nn.Module):
 953            self.main_module = start_module
 954        else:
 955            print("start_module must be nn.Module: %s" % name)
 956            print(type(start_module))
 957            print(start_module)
 958            sys.exit(-1)
 959        self.name = name
 960
 961        self.type = "tracked_module"
 962        set_tracked_params(self.main_module)
 963        if GPA.pc.get_verbose():
 964            print(
 965                f"tracking a module {self.name} with main type {type(self.main_module)}"
 966            )
 967            print(start_module)
 968        GPA.pai_tracker.add_tracked_neuron_module(self)
 969        if GPA.pc.get_perforated_backpropagation():
 970            MPB.set_neuron_parameters(self.main_module)
 971
 972    def __getattr__(self, name):
 973        """Get member variables from the main module.
 974
 975        Parameters
 976        ----------
 977        name : str
 978            The name of the variable to retrieve.
 979        Returns
 980        -------
 981        The requested variable.
 982
 983        Notes
 984        -----
 985        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
 986        If it fails, it tries to get the attribute from the wrapped main_module.
 987        This allows seamless access to the main module's attributes without modifying original code.
 988        """
 989        try:
 990            return super().__getattr__(name)
 991        except AttributeError:
 992            return getattr(self.main_module, name)
 993
 994    def __getitem__(self, index):
 995        """Support indexing operations on the main module.
 996
 997        Parameters
 998        ----------
 999        index : int or slice
1000            The index or slice to retrieve.
1001
1002        Returns
1003        -------
1004        The indexed item from the main module.
1005        """
1006        return self.main_module[index]
1007
1008    def set_mode(self, mode):
1009        """Set mode for tracked module.
1010
1011        Parameters
1012        ----------
1013        mode : str
1014            The mode to set. Either "n" for neuron training or "p" for pai-dendrite training.
1015
1016        Returns
1017        -------
1018        bool
1019            True.
1020
1021        Notes
1022        -----
1023        This function does not change any behavior since this is a tracked module.
1024        """
1025
1026        if GPA.pc.get_verbose():
1027            print(f"{self.name} calling set mode {mode}")
1028        return True
1029
1030    def forward(self, *args, **kwargs):
1031        """Forward pass for tracked module.
1032
1033        Parameters
1034        ----------
1035        *args : tuple
1036            Positional arguments for the forward pass.
1037        **kwargs : dict
1038            Keyword arguments for the forward pass.
1039
1040        Returns
1041        -------
1042        Any
1043            The output of the module
1044
1045        Notes
1046        -----
1047            The output of this forward function will have the same format as the output
1048            of the original module
1049        """
1050        return self.main_module(*args, **kwargs)
1051
1052    def __str__(self):
1053        """String representation of the module.
1054
1055        Parameters
1056        ----------
1057        None
1058
1059        Returns
1060        -------
1061        str
1062            String representation of the module.
1063
1064        Notes
1065        -----
1066        Setting for verbose changes level of details in the string output.
1067        """
1068
1069        if GPA.pc.get_verbose():
1070            total_string = self.main_module.__str__()
1071            total_string = "PAITrackedModule(" + total_string + ")"
1072            return total_string
1073        else:
1074            total_string = self.main_module.__str__()
1075            total_string = "PAITrackedModule(" + total_string + ")"
1076            return total_string
1077
1078    def __repr__(self):
1079        """Representation of the module."""
1080        return self.__str__()
1081
1082
1083def init_params(module, neuron_main_module):
1084    """Randomize weights after duplicating the main module for the next set of dendrites.
1085
1086    Parameters
1087    ----------
1088    module : nn.Module
1089        The new dendrite module to initialize.
1090    neuron_main_module : nn.Module
1091        The main module of the neuron for potential weight scaling.
1092
1093
1094    Returns
1095    -------
1096    None
1097        This function does not return a value.
1098    """
1099    for param in module.parameters():
1100        if param.dtype == torch.uint8:
1101            param.data = torch.randint(0, 256, param.size(), dtype=torch.uint8)
1102        else:
1103            # If factoring in the main modules weights multiply the randn()
1104            #  by the average abs value of the main modules weights
1105            if GPA.pc.get_candidate_weight_init_by_main():
1106                main_module_abs = 0
1107                total_main_params = 0
1108                for main_param in neuron_main_module.parameters():
1109                    main_module_abs += main_param.abs().sum().item()
1110                    total_main_params += main_param.numel()
1111                if total_main_params > 0:
1112                    main_module_abs /= total_main_params
1113                else:
1114                    main_module_abs = 1.0
1115                multiplier = main_module_abs
1116            else:
1117                multiplier = 1.0
1118            param.data = (
1119                torch.randn(param.size(), dtype=param.dtype)
1120                * GPA.pc.get_candidate_weight_initialization_multiplier()
1121                * multiplier
1122            )
1123
1124
1125class PAIDendriteModule(nn.Module):
1126    """Module containing all dendrites modules added to the neuron module."""
1127
1128    def __init__(
1129        self,
1130        initial_module,
1131        activation_function_value=0.3,
1132        name="no_name_given",
1133        output_dimensions=None,
1134    ):
1135        """Initialize PAINeuronModule.
1136
1137        This function sets up the dendrite module to create candidate and permanent
1138        dendrite modules based on the initial_module provided.
1139
1140        Parameters
1141        ----------
1142        initial_module : nn.Module
1143            The module to copy.
1144        activation_function_value : float, optional
1145            A value associated with the activation function, by default 0.3.
1146        name : str
1147            The name of the neuron module.
1148        output_dimensions : vector, optional
1149            The dimensions of the input vector
1150        """
1151        super(PAIDendriteModule, self).__init__()
1152
1153        if output_dimensions is None:
1154            output_dimensions = []
1155
1156        self.layers = nn.ModuleList([])
1157        self.processors = []
1158        self.candidate_processors = []
1159        self.num_dendrites = 0
1160        self._create_dendrite_fn = None
1161        # Number of dendrite cycles performed
1162        self.register_buffer(
1163            "num_cycles",
1164            torch.zeros(1, device=GPA.pc.get_device(), dtype=GPA.pc.get_d_type()),
1165        )
1166        self.mode = "n"
1167        self.name = name
1168        # Create a copy of the parent module so you don't have a pointer to the real one which causes save errors
1169        self.parent_module = UPA.deep_copy_pai(initial_module)
1170        if GPA.pc.get_perforated_backpropagation():
1171            MPB.set_ignored_parameters(self.parent_module)
1172        # Setup the input dimensions and node index for combining dendrite outputs
1173        if GPA.pc.get_perforated_backpropagation():
1174            MPB.create_extra_tensors(self)
1175        if output_dimensions == []:
1176            self.register_buffer(
1177                "this_output_dimensions", torch.tensor(GPA.pc.get_output_dimensions())
1178            )
1179        else:
1180            self.register_buffer(
1181                "this_output_dimensions", output_dimensions.detach().clone()
1182            )
1183        if (self.this_output_dimensions == 0).sum() != 1:
1184            print(f"1 need exactly one 0 in the input dimensions: {self.name}")
1185            print(self.this_output_dimensions)
1186            sys.exit(-1)
1187        self.register_buffer(
1188            "this_node_index", torch.tensor(GPA.pc.get_output_dimensions().index(0))
1189        )
1190
1191        # Initialize dendrite to dendrite connections
1192        self.dendrites_to_candidates = nn.ParameterList()
1193        self.dendrites_to_dendrites = nn.ParameterList()
1194
1195        # Store an activation function value if required
1196        self.activation_function_value = activation_function_value
1197        self.dendrite_values = nn.ModuleList([])
1198        for j in range(0, GPA.pc.get_global_candidates()):
1199            if GPA.pc.get_verbose():
1200                print(f"creating dendrite Values for {self.name}")
1201            self.dendrite_values.append(
1202                DendriteValueTracker(
1203                    False,
1204                    self.activation_function_value,
1205                    self.name,
1206                    self.this_output_dimensions,
1207                )
1208            )
1209        if GPA.pc.get_perforated_backpropagation():
1210            self.apply_pb_grads = MPB.apply_pb_grads.__get__(self, type(self))
1211            self.apply_pb_zero = MPB.apply_pb_zero.__get__(self, type(self))
1212
1213    def __getstate__(self):
1214        """Tell pickle what to save when this object is serialized (e.g. torch.save).
1215
1216        apply_pb_grads and apply_pb_zero are bound methods of functions defined
1217        in modules_pbp and cannot be pickled.  Strip them out; __setstate__ will
1218        re-attach them after loading.
1219        """
1220        import types
1221
1222        pickle_safe_state = {}
1223        for attr_name, attr_value in self.__dict__.items():
1224            if not isinstance(attr_value, types.MethodType):
1225                pickle_safe_state[attr_name] = attr_value
1226
1227        return pickle_safe_state
1228
1229    def __setstate__(self, saved_state):
1230        """Restore this object from a pickled state (e.g. torch.load).
1231
1232        Restores all normal attributes, then re-attaches apply_pb_grads and
1233        apply_pb_zero if perforated backpropagation is enabled.
1234        """
1235        self.__dict__.update(saved_state)
1236
1237        # Re-attach the PBP bound methods that were stripped by __getstate__.
1238        # dendrite_loss_fn being present on the saved state means PBP was active
1239        # when the checkpoint was saved.
1240        if "dendrite_loss_fn" in saved_state:
1241            import perforatedbp.modules_pbp as MPB
1242            self.apply_pb_grads = MPB.apply_pb_grads.__get__(self, type(self))
1243            self.apply_pb_zero = MPB.apply_pb_zero.__get__(self, type(self))
1244
1245    def create_dendrite(self, parent_module):
1246        """Create a dendrite module from the parent module.
1247
1248        Override this function via set_create_dendrite to control how the dendrite
1249        module is created (e.g. to avoid a deep copy).
1250
1251        Parameters
1252        ----------
1253        parent_module : nn.Module
1254            The module to create a dendrite from.
1255
1256        Returns
1257        -------
1258        nn.Module
1259            The new dendrite module.
1260        """
1261        if self._create_dendrite_fn is not None:
1262            return self._create_dendrite_fn(parent_module)
1263        return UPA.deep_copy_pai(parent_module)
1264
1265    def set_create_dendrite(self, fn):
1266        """Set a custom function for creating dendrite modules.
1267
1268        Call this on a PAIDendriteModule instance to override how dendrites are
1269        created from the parent module. The function receives the parent module
1270        and must return a new nn.Module.
1271
1272        Parameters
1273        ----------
1274        fn : callable
1275            A function with signature ``fn(parent_module) -> nn.Module``.
1276
1277        Returns
1278        -------
1279        None
1280        """
1281        self._create_dendrite_fn = fn
1282
1283
1284    def set_this_output_dimensions(self, new_output_dimensions):
1285        """Set input dimensions for dendrite module.
1286
1287        Signals to this DendriteModule that its input dimensions are different
1288        than the global default.
1289
1290        Parameters
1291        ----------
1292        new_output_dimensions : list
1293            A list or tensor specifying the new input dimensions.
1294        Returns
1295        -------
1296        None
1297
1298        """
1299
1300        if type(new_output_dimensions) is list:
1301            new_output_dimensions = torch.tensor(new_output_dimensions)
1302        delattr(self, "this_output_dimensions")
1303        self.register_buffer(
1304            "this_output_dimensions", new_output_dimensions.detach().clone()
1305        )
1306        if (new_output_dimensions == 0).sum() != 1:
1307            print(f"2 Need exactly one 0 in the input dimensions: {self.name}")
1308            print(new_output_dimensions)
1309            sys.exit(-1)
1310        self.this_node_index.copy_(
1311            (new_output_dimensions == 0).nonzero(as_tuple=True)[0][0]
1312        )
1313        for j in range(0, GPA.pc.get_global_candidates()):
1314            self.dendrite_values[j].set_this_output_dimensions(new_output_dimensions)
1315
1316    def create_new_dendrite_module(self, neuron_main_module):
1317        """Add a new set of dendrites.
1318
1319        Parameters
1320        ----------
1321        neuron_main_module : Any
1322            PyTorch module to be used for dendritic learning.
1323                Typically a copy of the original neuron module.
1324
1325        Returns
1326        -------
1327        None
1328            This function does not return a value.
1329        """
1330        # Candidate module
1331        self.candidate_module = nn.ModuleList([])
1332        # Copy that is unused for open source version
1333        self.best_candidate_module = nn.ModuleList([])
1334        if GPA.pc.get_verbose():
1335            print(self.name)
1336            print("Setting candidate processors")
1337        self.candidate_processors = []
1338        with torch.no_grad():
1339            for i in range(0, GPA.pc.get_global_candidates()):
1340
1341                new_module = self.create_dendrite(self.parent_module)
1342                init_params(new_module, neuron_main_module)
1343                self.candidate_module.append(new_module)
1344                self.best_candidate_module.append(self.create_dendrite(new_module))
1345                if type(self.parent_module) in GPA.pc.get_modules_with_processing():
1346                    module_index = GPA.pc.get_modules_with_processing().index(
1347                        type(self.parent_module)
1348                    )
1349                    self.candidate_processors.append(
1350                        GPA.pc.get_modules_processing_classes()[module_index]()
1351                    )
1352                elif (
1353                    type(self.parent_module).__name__
1354                    in GPA.pc.get_module_names_with_processing()
1355                ):
1356                    module_index = GPA.pc.get_module_names_with_processing().index(
1357                        type(self.parent_module).__name__
1358                    )
1359                    self.candidate_processors.append(
1360                        GPA.pc.get_module_by_name_processing_classes()[module_index]()
1361                    )
1362                if GPA.pc.get_perforated_backpropagation():
1363                    MPB.set_candidate_parameters(self.candidate_module[i])
1364                    MPB.set_ignored_parameters(self.best_candidate_module[i])
1365
1366        for i in range(0, GPA.pc.get_global_candidates()):
1367            self.candidate_module[i].to(GPA.pc.get_device())
1368            self.best_candidate_module[i].to(GPA.pc.get_device())
1369
1370        # Reset the dendrite_values objects
1371        for j in range(0, GPA.pc.get_global_candidates()):
1372            self.dendrite_values[j].reinitialize_for_pai()
1373
1374        # If there are already dendrites initialize the dendrite to dendrite connections
1375        if self.num_dendrites > 0:
1376            self.dendrites_to_candidates = nn.ParameterList()
1377            for j in range(0, GPA.pc.get_global_candidates()):
1378                self.dendrites_to_candidates.append(
1379                    nn.Parameter(
1380                        torch.zeros(
1381                            (self.num_dendrites, self.out_channels),
1382                            device=GPA.pc.get_device(),
1383                            dtype=GPA.pc.get_d_type(),
1384                        ),
1385                        requires_grad=True,
1386                    )
1387                )
1388                if GPA.pc.get_perforated_backpropagation():
1389                    MPB.init_candidates(self, j)
1390            if GPA.pc.get_perforated_backpropagation():
1391                MPB.set_candidate_parameters(self.dendrites_to_candidates)
1392            # Initialize best_dendrites_to_candidates_saved to snapshot peak-correlation weights at epoch boundaries
1393            self.best_dendrites_to_candidates_saved = []
1394            for j in range(0, GPA.pc.get_global_candidates()):
1395                self.best_dendrites_to_candidates_saved.append(
1396                    torch.zeros(
1397                        (self.num_dendrites, self.out_channels),
1398                        device=GPA.pc.get_device(),
1399                        dtype=GPA.pc.get_d_type(),
1400                    )
1401                )
1402
1403    def clear_processors(self):
1404        """Clear processors.
1405
1406        Parameters
1407        ----------
1408        None
1409
1410        Returns
1411        -------
1412        None
1413            This function does not return a value.
1414        """
1415        for processor in self.processors:
1416            if not processor:
1417                continue
1418            else:
1419                processor.clear_processor()
1420        for processor in self.candidate_processors:
1421            if not processor:
1422                continue
1423            else:
1424                processor.clear_processor()
1425
1426    def set_mode(self, mode):
1427        """Perform actions when switching between neuron and dendrite training.
1428
1429        Parameters
1430        ----------
1431        mode : str
1432            The mode to set. Either "n" for neuron training or "p" for pai-dendrite training.
1433
1434        Returns
1435        -------
1436        None
1437        """
1438
1439        self.mode = mode
1440        self.num_cycles += 1
1441        if GPA.pc.get_verbose():
1442            print(f"PAI calling set mode {mode} : {self.num_cycles}")
1443        if not GPA.pc.get_silent():
1444            print(f"Module {self.name} calling set mode {mode} : {self.num_cycles}")
1445        # When switching back to neuron training mode convert candidates modules into accepted modules
1446        if mode == "n":
1447            if GPA.pc.get_verbose():
1448                print("So calling all the things to add to modules")
1449            # Copy weights/bias from correct candidates
1450            if self.num_dendrites == 1:
1451                self.dendrites_to_dendrites = nn.ParameterList()
1452                self.dendrites_to_dendrites.append(torch.tensor([]))
1453            if self.num_dendrites >= 1:
1454                self.dendrites_to_dendrites.append(
1455                    torch.nn.Parameter(
1456                        torch.zeros(
1457                            [self.num_dendrites, self.out_channels],
1458                            device=GPA.pc.get_device(),
1459                            dtype=GPA.pc.get_d_type(),
1460                        ),
1461                        # Grad is true if not pb or if pb and dendrite_update_mode is true
1462                        requires_grad=(not GPA.pc.get_perforated_backpropagation())
1463                        or GPA.pc.get_dendrite_update_mode(),
1464                    )
1465                )
1466            with torch.no_grad():
1467                if GPA.pc.get_global_candidates() > 1:
1468                    print(
1469                        "This was a flag that will be needed if using multiple candidates. "
1470                        "It's not set up yet but nice work finding it."
1471                    )
1472                    print(
1473                        "Note: with multiple candidates, best-score ranking in new_best() uses "
1474                        "unnormalized covariance (prev_dendrite_candidate_correlation) rather than "
1475                        "the normalized correlation coefficient. Candidates with larger output "
1476                        "magnitude will be favored regardless of true correlation quality. "
1477                        "Fix by tracking running sigma_V and sigma_E and dividing in new_best()."
1478                    )
1479                    pdb.set_trace()
1480                plane_max_index = 0
1481                self.layers.append(
1482                    UPA.deep_copy_pai(self.best_candidate_module[plane_max_index])
1483                )
1484                self.layers[self.num_dendrites].to(GPA.pc.get_device())
1485                if self.num_dendrites > 0:
1486                    self.dendrites_to_dendrites[self.num_dendrites].copy_(
1487                        self.best_dendrites_to_candidates_saved[plane_max_index]
1488                    )
1489                if type(self.parent_module) in GPA.pc.get_modules_with_processing():
1490                    self.processors.append(self.candidate_processors[plane_max_index])
1491                if (
1492                    type(self.parent_module).__name__
1493                    in GPA.pc.get_module_names_with_processing()
1494                ):
1495                    self.processors.append(self.candidate_processors[plane_max_index])
1496            if GPA.pc.get_perforated_backpropagation():
1497                MPB.set_pb_mode(self, mode)
1498            del self.candidate_module, self.best_candidate_module
1499
1500            self.num_dendrites += 1
1501            if GPA.pc.get_perforated_backpropagation():
1502                MPB.set_dendrite_parameters(self.dendrites_to_dendrites)
1503                MPB.set_dendrite_parameters(self.layers)
1504
1505    def forward(self, *args, **kwargs):
1506        """Forward pass for dendrite module.
1507
1508        Parameters
1509        ----------
1510        *args : tuple
1511            Positional arguments for the forward pass.
1512        **kwargs : dict
1513            Keyword arguments for the forward pass.
1514
1515        Returns
1516        -------
1517        Any
1518            The output of the module after processing through the neuron and dendrite modules.
1519        Any
1520            Remaining outputs are only used for Perforated Backpropagation.
1521        Any
1522            Remaining outputs are only used for Perforated Backpropagation.
1523        Any
1524            Remaining outputs are only used for Perforated Backpropagation.
1525
1526        Notes
1527        -----
1528        If using Perforated Backpropagation, the additional outputs will be moved around in
1529        this code but left unused and only passed into separate PB functions.
1530        """
1531
1532        outs = {}
1533
1534        # For all modules apply processors, call the modules, then apply post processors
1535        args2, kwargs2 = args, kwargs
1536        for c in range(0, self.num_dendrites):
1537            if GPA.pc.get_perforated_backpropagation():
1538                args2, kwargs2 = MPB.preprocess_pb(*args, **kwargs)
1539            if self.processors != []:
1540                try:
1541                    args2, kwargs2 = self.processors[c].pre_d(*args2, **kwargs2)
1542                except Exception as e:
1543                    traceback.print_exc(limit=None, chain=True)
1544                    print(f"Your pre_d processor for {self.name} caused this error")
1545                    print(
1546                        f"You must check how this is defined and ensure that it is properly"
1547                    )
1548                    print(
1549                        f"accepting inputs to the PAIModule and returning what will then be"
1550                    )
1551                    print(f"the input to the dendrite module")
1552                    sys.exit()
1553            out_values = self.layers[c](*args2, **kwargs2)
1554            if self.processors != []:
1555                try:
1556                    outs[c] = self.processors[c].post_d(out_values)
1557                except Exception as e:
1558                    traceback.print_exc(limit=None, chain=True)
1559                    print(f"Your post_d processor for {self.name} caused this error")
1560                    print(
1561                        f"You must check how this is defined and ensure that it is properly"
1562                    )
1563                    print(
1564                        f"accepting outputs from the dendrite module and returning the"
1565                    )
1566                    print(
1567                        f"single tensor to be combined with the neurons output tensor"
1568                    )
1569                    sys.exit()
1570            else:
1571                outs[c] = out_values
1572
1573        # Create dendrite outputs
1574        # Each dendrite has input from previously created dendrites
1575        # So activation is added before the nonlinearity is called
1576        view_tuple = []
1577        for out_index in range(0, self.num_dendrites):
1578            current_out = outs[out_index]
1579            view_tuple = []
1580            for dim in range(len(current_out.shape)):
1581                if dim == self.this_node_index:
1582                    view_tuple.append(-1)
1583                    continue
1584                view_tuple.append(1)
1585
1586            for in_index in range(0, out_index):
1587                if view_tuple == [
1588                    1
1589                ]:  # This is only the case when passing a single datapoint rather than a batch
1590                    current_out = (
1591                        current_out
1592                        + self.dendrites_to_dendrites[out_index][in_index, :].to(
1593                            current_out.device
1594                        )
1595                        * outs[in_index]
1596                    )
1597                else:
1598                    current_out = (
1599                        current_out
1600                        + self.dendrites_to_dendrites[out_index][in_index, :]
1601                        .view(view_tuple)
1602                        .to(current_out.device)
1603                        * outs[in_index]
1604                    )
1605            outs[out_index] = GPA.pc.get_pai_forward_function()(current_out)
1606        # Return a dict which has all dendritic outputs after the activation functions were called
1607        if GPA.pc.get_perforated_backpropagation():
1608            candidate_outs, candidate_nonlinear_outs, candidate_non_zeroed = (
1609                MPB.forward_candidates(self, view_tuple, outs, *args2, **kwargs2)
1610            )
1611        else:
1612            candidate_outs, candidate_nonlinear_outs, candidate_non_zeroed = (
1613                {},
1614                {},
1615                {},
1616            )
1617        return outs, candidate_outs, candidate_nonlinear_outs, candidate_non_zeroed
1618
1619
1620class DendriteValueTracker(nn.Module):
1621    """Tracker object that maintains certain values for each set of dendrites."""
1622
1623    def __init__(
1624        self,
1625        initialized,
1626        activation_function_value,
1627        name,
1628        output_dimensions,
1629        out_channels=-1,
1630    ):
1631        """Initialize DendriteValueTracker.
1632
1633        This function sets up the value tracker to maintain statistics and values
1634        for each set of dendrites.
1635
1636        Parameters
1637        ----------
1638        initialized : int
1639            Whether the dendrite has been initialized (1) or not (0).
1640        activation_function_value : float
1641            A value associated with the activation function.
1642        name : str
1643            The name of the associated neuron module.
1644        output_dimensions : vector
1645            The dimensions of the input vector.
1646        out_channels : int
1647            The number of output channels
1648        """
1649        super(DendriteValueTracker, self).__init__()
1650
1651        self.layer_name = name
1652        for val_name in DENDRITE_INIT_VALUES:
1653            self.register_buffer(
1654                val_name,
1655                torch.zeros(1, device=GPA.pc.get_device(), dtype=GPA.pc.get_d_type()),
1656            )
1657        self.initialized[0] = initialized
1658        self.activation_function_value = activation_function_value
1659        self.register_buffer(
1660            "this_output_dimensions", output_dimensions.clone().detach()
1661        )
1662        if (self.this_output_dimensions == 0).sum() != 1:
1663            print(f"3 need exactly one 0 in the input dimensions: {self.layer_name}")
1664            print(self.this_output_dimensions)
1665            sys.exit(-1)
1666        self.register_buffer(
1667            "this_node_index", (output_dimensions == 0).nonzero(as_tuple=True)[0]
1668        )
1669        if out_channels != -1:
1670            ndim = len(output_dimensions)
1671            init_shape = [1] * ndim
1672            init_shape[(output_dimensions == 0).nonzero(as_tuple=True)[0].item()] = out_channels
1673            self.setup_arrays(init_shape)
1674        else:
1675            self.out_channels = -1
1676
1677    def print(self):
1678        """Print value tracker information.
1679
1680        Parameters
1681        ----------
1682        None
1683
1684        Returns
1685        -------
1686        None
1687            This function does not return a value.
1688        """
1689        total_string = "Value Tracker:"
1690        for val_name in DENDRITE_INIT_VALUES:
1691            total_string += f"\t{val_name}:\n\t\t"
1692            total_string += getattr(self, val_name).__repr__()
1693            total_string += "\n"
1694        for val_name in get_DENDRITE_TENSOR_VALUES():
1695            if getattr(self, val_name, None) is not None:
1696                total_string += f"\t{val_name}:\n\t\t"
1697                total_string += getattr(self, val_name).__repr__()
1698                total_string += "\n"
1699        print(total_string)
1700
1701    def set_this_output_dimensions(self, new_output_dimensions):
1702        """Set input dimensions for value tracker
1703
1704        Signals to this DendriteValueTracker that its input dimensions are different
1705        than the global default.
1706
1707        Parameters
1708        ----------
1709        new_output_dimensions : list
1710            A list or tensor specifying the new input dimensions.
1711        Returns
1712        -------
1713        None
1714
1715        """
1716        if type(new_output_dimensions) is list:
1717            new_output_dimensions = torch.tensor(new_output_dimensions)
1718        delattr(self, "this_output_dimensions")
1719        self.register_buffer(
1720            "this_output_dimensions", new_output_dimensions.detach().clone()
1721        )
1722        if (new_output_dimensions == 0).sum() != 1:
1723            print(f"4 need exactly one 0 in the input dimensions: {self.layer_name}")
1724            print(new_output_dimensions)
1725            sys.exit(-1)
1726        self.this_node_index.copy_(
1727            (new_output_dimensions == 0).nonzero(as_tuple=True)[0][0]
1728        )
1729
1730    def set_out_channels(self, shape_values):
1731        """Set output channels based on shape values and saved node index
1732
1733        Parameters
1734        ----------
1735        shape_values : list or torch.Size
1736            A list or tensor specifying the shape values.
1737
1738        Returns
1739        -------
1740        None
1741        """
1742        if type(shape_values) == torch.Size:
1743            self.out_channels = int(shape_values[self.this_node_index])
1744        else:
1745            self.out_channels = int(shape_values[self.this_node_index].item())
1746
1747    def setup_arrays(self, storage_shape):
1748        """Setup arrays for value tracker.
1749
1750        Parameters
1751        ----------
1752        storage_shape : list
1753            Shape for the tracking tensors: 1 at every dim except the channel
1754            dim (this_node_index), which holds out_channels.  E.g. [1, 5] for
1755            a linear layer with 5 outputs.
1756        Returns
1757        -------
1758        None
1759
1760        """
1761        # storage_shape is a list with 1 at every dim except the channel dim.
1762        # Passed in directly from filter_backward so it is always derived from
1763        # the live gradient — not from out_channels, which is not saved/loaded.
1764        self.out_channels = storage_shape[self.this_node_index.item()]
1765        self.register_buffer(
1766            "dendrite_storage_shape",
1767            torch.tensor(storage_shape, dtype=torch.long, device=GPA.pc.get_device()),
1768        )
1769        for val_name in get_DENDRITE_TENSOR_VALUES():
1770            self.register_buffer(
1771                val_name,
1772                torch.zeros(
1773                    storage_shape, device=GPA.pc.get_device(), dtype=GPA.pc.get_d_type()
1774                ),
1775            )
1776
1777        for name in get_VALUE_TRACKER_ARRAYS():
1778            setattr(self, name, {})
1779            count = 1
1780            if torch.cuda.device_count() > count:
1781                count = torch.cuda.device_count()
1782            for i in range(count):
1783                getattr(self, name)[i] = []
1784        for val_name in get_DENDRITE_SINGLE_VALUES():
1785            self.register_buffer(
1786                val_name,
1787                torch.zeros(1, device=GPA.pc.get_device(), dtype=GPA.pc.get_d_type()),
1788            )
1789
1790        # If fixed_input_sizes is enabled, register cache buffers now that the
1791        # final ndim is known (output_dimensions may have been corrected after
1792        # __init__ for Linear layers). Mirrors the pattern of this_output_dimensions
1793        # — registered as buffers so they are saved and loaded automatically.
1794        if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_fixed_input_sizes():
1795            ndim = len(storage_shape)
1796            if not hasattr(self, 'math_tuple_cache'):
1797                self.register_buffer(
1798                    "math_tuple_cache",
1799                    torch.zeros(ndim, dtype=torch.long, device=GPA.pc.get_device()),
1800                )
1801                self.register_buffer(
1802                    "view_tuple_cache",
1803                    torch.zeros(ndim, dtype=torch.long, device=GPA.pc.get_device()),
1804                )
1805                self.register_buffer(
1806                    "full_mult_cache",
1807                    torch.zeros(1, dtype=torch.long, device=GPA.pc.get_device()),
1808            )
1809
1810    def reinitialize_for_pai(self):
1811        """Reinitialize value tracker to add the next set of dendrites
1812
1813        Parameters
1814        ----------
1815        None
1816
1817        Returns
1818        -------
1819        None
1820            This function does not return a value.
1821        """
1822
1823        if self.out_channels == -1:
1824            print("You have a perforated module that was never initialized")
1825            print("This likely means it is not being added to the autograd graph")
1826            print("Check your forward function that it is actually being used")
1827            print("If its not you should really delete it, but you can also add")
1828            print(self.layer_name)
1829            print("with:")
1830            print("GPA.pc.append_module_ids_to_track(['" + self.layer_name + "'])")
1831            print("This can also happen while testing_dendrite_capacity if you")
1832            print(
1833                "run a validation cycle and try to add Dendrites before doing any training.\n"
1834            )
1835            pdb.set_trace()
1836
1837        self.initialized[0] = 0
1838        if GPA.pc.get_perforated_backpropagation():
1839            MPB.reinitialize_for_pb(self)
1840        else:
1841            for val_name in get_DENDRITE_REINIT_VALUES():
1842                setattr(self, val_name, getattr(self, val_name) * 0)