perforatedai.modules_perforatedai

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