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