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