ngclearn 3.2.2__tar.gz → 3.2.4__tar.gz

{ngclearn-3.2.2/ngclearn.egg-info → ngclearn-3.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ngclearn
-Version: 3.2.2
+Version: 3.2.4
 Summary: Simulation software for building and analyzing computational neuroscience models, brain-inspired computing systems, and NeuroAI agents.
 Author-email: Alexander Ororbia <ago@cs.rit.edu>, William Gebhardt <wdg1351@rit.edu>
 License-Expression: BSD-3-Clause
@@ -40,7 +40,8 @@ Dynamic: license-file
 
 [![Python Version](https://img.shields.io/badge/python-3.10%20%7C%203.11-blue.svg)](https://www.python.org/downloads)[![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://GitHub.com/Naereen/StrapDown.js/graphs/commit-activity)[![Documentation Status](https://readthedocs.org/projects/ngc-learn/badge/?version=latest)](http://ngc-learn.readthedocs.io/en/latest/?badge=latest)[![DOI](https://zenodo.org/badge/483413212.svg)](https://zenodo.org/badge/latestdoi/483413212)
 
-<img src="docs/images/ngc-learn-logo.png" width="300">
+<!-- <img src="docs/images/ngc-learn-logo.png" width="300"> -->
+<img src="https://raw.githubusercontent.com/NACLab/ngc-learn/main/docs/images/ngc-learn-logo.png" width="300">
 
 <b>ngc-learn</b> is a Python library for building, simulating, and analyzing biophysical / neurobiological systems, spiking neuronal networks, predictive coding circuitry, and biomimetic (NeuroAI) agents that learn in a biologically-plausible manner. This simulation toolkit, meant to support computational neuroscience and brain-inspired computing research, is built on top of JAX and is distributed under the 3-Clause BSD license.
 
@@ -66,24 +67,22 @@ chosen as one of the <i>Editors' Highlights for Applied Physics and Mathematics<
 ### Dependencies
 
 ngc-learn requires:
-1) Python (>=3.10)
-2) NumPy (>=1.22.0)
-3) SciPy (>=1.7.0)
+1) Python (>=3.12)
+2) NumPy (>=2.5.2)
+3) SciPy (>=1.18.1)
 4) ngcsimlib (>=3.1.1), (visit official page <a href="https://github.com/NACLab/ngc-sim-lib">here</a>)
-5) JAX (>=0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
+5) JAX (>=0.11.1) (to enable GPU use, make sure to install one of the CUDA variants)
 <!--
-5) scikit-learn (>=1.3.1) if using `ngclearn.utils.density`
-6) matplotlib (>=3.4.3) if using `ngclearn.utils.viz`
 6) networkx  (>=2.6.3) (currently optional but required if using `ngclearn.utils.experimental.viz_utils`)
 7) pyviz (>=0.2.0) (currently optional but required if using `ngclearn.utils.experimental.viz_utils`)
 -->
 
 ---
-ngc-learn 3.2.1 and later require Python 3.10 or newer as well as ngcsimlib >=3.1.1. 
+ngc-learn 3.2.4 and later require Python 3.12 or newer as well as ngcsimlib >=3.1.1. 
 ngc-learn's plotting capabilities (routines within `ngclearn.utils.viz`) require
-Matplotlib (>=3.8.0) and imageio (>=2.31.5) and both plotting and density estimation
-tools (routines within ``ngclearn.utils.density``) will require Scikit-learn (>=0.24.2).
-Many of the tutorials will require Matplotlib (>=3.8.0), imageio (>=2.31.5), and Scikit-learn (>=0.24.2).
+Matplotlib (>=3.11.1) and imageio (>=2.37.4) and both plotting and density estimation
+tools (routines within ``ngclearn.utils.density``) will require Scikit-learn (>=1.9.0).
+Many of the tutorials will require Matplotlib (>=3.11.1), imageio (>=2.37.4), and Scikit-learn (>=1.9.0).
 
 <i>Note</i>: If you are working with Cuda 12 and want to use jax/jaxlib versions > 0.4.28, you might need to 
 check that you are working with the right version of Cudnn (e.g., `nvidia-cudnn-cu12==9.10.2.21`) to ensure 
@@ -133,7 +132,7 @@ Python 3.12.13 (main, MONTH  DAY YEAR, TIME) [GCC XX.X.X] on linux
 Type "help", "copyright", "credits" or "license" for more information.
 >>> import ngclearn
 >>> ngclearn.__version__
-'3.2.1'
+'3.2.4'
 ```
 
 <i>Note:</i> For access to the previous Tensorflow-2 version of ngc-learn (of
@@ -180,7 +179,7 @@ $ python install -e .
 </pre>
 
 **Version:**<br>
-3.2.2
+3.2.4
 
 Author:
 Alexander G. Ororbia II<br>

{ngclearn-3.2.2 → ngclearn-3.2.4}/README.md

@@ -1,6 +1,7 @@
 [![Python Version](https://img.shields.io/badge/python-3.10%20%7C%203.11-blue.svg)](https://www.python.org/downloads)[![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://GitHub.com/Naereen/StrapDown.js/graphs/commit-activity)[![Documentation Status](https://readthedocs.org/projects/ngc-learn/badge/?version=latest)](http://ngc-learn.readthedocs.io/en/latest/?badge=latest)[![DOI](https://zenodo.org/badge/483413212.svg)](https://zenodo.org/badge/latestdoi/483413212)
 
-<img src="docs/images/ngc-learn-logo.png" width="300">
+<!-- <img src="docs/images/ngc-learn-logo.png" width="300"> -->
+<img src="https://raw.githubusercontent.com/NACLab/ngc-learn/main/docs/images/ngc-learn-logo.png" width="300">
 
 <b>ngc-learn</b> is a Python library for building, simulating, and analyzing biophysical / neurobiological systems, spiking neuronal networks, predictive coding circuitry, and biomimetic (NeuroAI) agents that learn in a biologically-plausible manner. This simulation toolkit, meant to support computational neuroscience and brain-inspired computing research, is built on top of JAX and is distributed under the 3-Clause BSD license.
 
@@ -26,24 +27,22 @@ chosen as one of the <i>Editors' Highlights for Applied Physics and Mathematics<
 ### Dependencies
 
 ngc-learn requires:
-1) Python (>=3.10)
-2) NumPy (>=1.22.0)
-3) SciPy (>=1.7.0)
+1) Python (>=3.12)
+2) NumPy (>=2.5.2)
+3) SciPy (>=1.18.1)
 4) ngcsimlib (>=3.1.1), (visit official page <a href="https://github.com/NACLab/ngc-sim-lib">here</a>)
-5) JAX (>=0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
+5) JAX (>=0.11.1) (to enable GPU use, make sure to install one of the CUDA variants)
 <!--
-5) scikit-learn (>=1.3.1) if using `ngclearn.utils.density`
-6) matplotlib (>=3.4.3) if using `ngclearn.utils.viz`
 6) networkx  (>=2.6.3) (currently optional but required if using `ngclearn.utils.experimental.viz_utils`)
 7) pyviz (>=0.2.0) (currently optional but required if using `ngclearn.utils.experimental.viz_utils`)
 -->
 
 ---
-ngc-learn 3.2.1 and later require Python 3.10 or newer as well as ngcsimlib >=3.1.1. 
+ngc-learn 3.2.4 and later require Python 3.12 or newer as well as ngcsimlib >=3.1.1. 
 ngc-learn's plotting capabilities (routines within `ngclearn.utils.viz`) require
-Matplotlib (>=3.8.0) and imageio (>=2.31.5) and both plotting and density estimation
-tools (routines within ``ngclearn.utils.density``) will require Scikit-learn (>=0.24.2).
-Many of the tutorials will require Matplotlib (>=3.8.0), imageio (>=2.31.5), and Scikit-learn (>=0.24.2).
+Matplotlib (>=3.11.1) and imageio (>=2.37.4) and both plotting and density estimation
+tools (routines within ``ngclearn.utils.density``) will require Scikit-learn (>=1.9.0).
+Many of the tutorials will require Matplotlib (>=3.11.1), imageio (>=2.37.4), and Scikit-learn (>=1.9.0).
 
 <i>Note</i>: If you are working with Cuda 12 and want to use jax/jaxlib versions > 0.4.28, you might need to 
 check that you are working with the right version of Cudnn (e.g., `nvidia-cudnn-cu12==9.10.2.21`) to ensure 
@@ -93,7 +92,7 @@ Python 3.12.13 (main, MONTH  DAY YEAR, TIME) [GCC XX.X.X] on linux
 Type "help", "copyright", "credits" or "license" for more information.
 >>> import ngclearn
 >>> ngclearn.__version__
-'3.2.1'
+'3.2.4'
 ```
 
 <i>Note:</i> For access to the previous Tensorflow-2 version of ngc-learn (of
@@ -140,7 +139,7 @@ $ python install -e .
 </pre>
 
 **Version:**<br>
-3.2.2
+3.2.4
 
 Author:
 Alexander G. Ororbia II<br>

{ngclearn-3.2.2 → ngclearn-3.2.4}/ngclearn/components/__init__.py

@@ -73,5 +73,13 @@ from .synapses.modulated.REINFORCESynapse import REINFORCESynapse
 ## point to patched component types
 from .synapses.patched.patchedSynapse import PatchedSynapse
 from .synapses.patched.staticPatchedSynapse import StaticPatchedSynapse
+from .synapses.patched.scorePatchedSynapse import ScorePatchedSynapse
 from .synapses.patched.hebbianPatchedSynapse import HebbianPatchedSynapse
 
+## point to attention patched synaptic components
+from .synapses.attention.attentionPatchedSynapse import AttentionPatchedSynapse
+from .synapses.attention.hebbianAttentionPatchedSynapse import HebbianAttentionPatchedSynapse
+from .synapses.attention.scoreAttentionPatchedSynapse import ScoreAttentionPatchedSynapse
+
+
+

{ngclearn-3.2.2 → ngclearn-3.2.4}/ngclearn/components/neurons/graded/rateCell.py

@@ -49,7 +49,6 @@ def _modulate(j, dfx_val):
     """
     return j * dfx_val
 
-# @partial(jit, static_argnames=["integType", "priorType"])
 def _run_cell(dt, j, j_td, z, tau_m, leak_gamma=0., integType=0, priorType=0):
     """
     Runs leaky rate-coded state dynamics one step in time.
@@ -166,7 +165,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
 
         use_lateral: determines wheter the RateCell uses lateral pressure during inference (Default: False)
 
-        n_patch:  the number of local patches in each layer (equivalent to n_sub_models in PatchedSynapse) (Default: 1)
+        n_patch:  the number of local patches in each layer (equivalent to n_patch in PatchedSynapse) (Default: 1)
 
         adaptive_lateral: if True, lateral synaptic is adaptive using HebbianPatchedSynapse component,
                           if False, lateral synaptic is a static using StaticPatchedSynapse component
@@ -265,7 +264,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
 
             if adaptive_lateral:
                 self.Wl = HebbianPatchedSynapse(f"Wl_{name}", shape=(n_units, n_units),
-                                                n_sub_models=n_patch,
+                                                n_patch=n_patch,
                                                 weight_init=dist.diag_offdiag(on_diag=exc, off_diag=inh),
                                                 eta=Wl_eta,
                                                 prior=Wl_prior,
@@ -275,7 +274,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
                                                 )
             else:
                 self.Wl = StaticPatchedSynapse(f"Wl_{name}", shape=(n_units, n_units),
-                                               n_sub_models=n_patch,
+                                               n_patch=n_patch,
                                                weight_init=dist.diag_offdiag(on_diag=exc, off_diag=inh),
                                                batch_size=batch_size,
                                                )
@@ -288,6 +287,8 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
         self.zF = Compartment(restVals, display_name="Transformed Rate Activity") # rate-coded output - activity
         self.j_td = Compartment(restVals, display_name="Modulatory Stimulus Current", units="mA") # top-down electrical current - pressure
         self.z = Compartment(restVals, display_name="Rate Activity", units="mA") # rate activity
+        self.mu_top = Compartment(restVals, display_name="Rate Activity", units="mA") # rate activity
+        self.mu_top_zF = Compartment(restVals, display_name="Transformed Rate Activity") # rate-coded output - activity
         ## lateral compartments
         if use_lateral:
             self.j_lat = Compartment(restVals, display_name="Lateral Stimulus Current", units="mA")  # lateral current
@@ -309,6 +310,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
         j = self.j.get()
         j_td = self.j_td.get()
         z = self.z.get()
+        mu_top = self.mu_top.get()
 
         if self.is_stateful:
             ## ════════════════════════════════════════════════════════════════
@@ -320,7 +322,6 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
             dfx_val = self.dfx(z)
             j = _modulate(j, dfx_val)
             j = j * self.resist_scale
-
             ## ══════════════════════════
             ## Calculate lateral pressure
             if self.use_lateral:
@@ -355,6 +356,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
             ## ════════════════════════════════════════════════════════════════
             ## run in "stateless" mode (when no membrane time constant provided)
             z = _run_cell_stateless(j + j_td)  ## (j+j_td) >> z
+            mu_top = _run_cell_stateless(mu_top)
 
         ## ═══════════════════════════════════════════════
         ## fx(z) >> zF post-activation function value(s)
@@ -363,11 +365,16 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
             zF_lat = self.fx_lat(z) * self.output_scale
 
         ## ════════════════════
+        mu_top_zF = self.fx(mu_top) * self.output_scale
+
+        ## ═══════════════════════════════════════════════
         ## Update compartments
         self.j.set(j)
         self.j_td.set(j_td)
         self.z.set(z)
         self.zF.set(zF)
+        self.mu_top.set(mu_top)
+        self.mu_top_zF.set(mu_top_zF)
         if self.use_lateral:
             self.j_lat.set(j_lat)
             self.zF_lat.set(zF_lat)
@@ -391,6 +398,8 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
         self.j_td.set(restVals)
         self.z.set(restVals)
         self.zF.set(restVals)
+        self.mu_top.set(restVals)
+        self.mu_top_zF.set(restVals)
 
         if self.use_lateral:
             self.Wl.reset()
@@ -417,13 +426,16 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
             "inputs":
                 {"j": "External input stimulus value(s)",
                  "j_td": "External top-down input stimulus value(s); these get "
-                         "multiplied by the derivative of f(x), i.e., df(x)"},
+                         "multiplied by the derivative of f(x), i.e., df(x)",
+                 "mu_top": "Top-down Prediction"},
             "states":
                 {"z": "Update to rate-coded continuous dynamics; value at time t"},
             "outputs":
                 {"zF": "Nonlinearity/function applied to rate-coded dynamics; f(z)",
                  "zF_lat": "(if use_lateral is True) zF for lateral fx_lat(z)",
-                 "j_lat": "(if use_lateral is True) lateral pressure"},
+                 "j_lat": "(if use_lateral is True) lateral pressure",
+                 "mu_top_zF": "Nonlinearity/function applied to the top-down prediction",
+                 },
         }
         hyperparams = {
             "n_units": "Number of neuronal cells to model in this layer",
@@ -434,7 +446,7 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
             "threshold": "What kind of iterative thresholding function to place over neuronal dynamics?",
             "integration_type": "Type of numerical integration to use for the cell dynamics",
             "use_lateral": "If lateral synaptic is active for inference",
-            "n_patch": "Number of block-diagonal local lateral synaptic patch (n_sub_models of `Wl`)",
+            "n_patch": "Number of block-diagonal local lateral synaptic patch (n_patch of `Wl`)",
             "adaptive_lateral": "Adaptive lateral synapses (HebbianPatchedSynapse) or static lateral synaptic (StaticPatchedSynapse)",
             "Wl_eta": "Learning rate for the lateral synapses",
             "exc_inh": "(exc, inh) lateral synapses initial values. On-diagonal takes exc>0 and off-diagonal takes inh<0",

{ngclearn-3.2.2 → ngclearn-3.2.4}/ngclearn/components/synapses/__init__.py

@@ -39,5 +39,12 @@ from .modulated.MSTDPETSynapse import MSTDPETSynapse
 ## patched synaptic components
 from .patched.patchedSynapse import PatchedSynapse
 from .patched.staticPatchedSynapse import StaticPatchedSynapse
+from .patched.scorePatchedSynapse import ScorePatchedSynapse
 from .patched.hebbianPatchedSynapse import HebbianPatchedSynapse
 
+## attention patched synaptic components
+from .attention.attentionPatchedSynapse import AttentionPatchedSynapse
+from .attention.hebbianAttentionPatchedSynapse import HebbianAttentionPatchedSynapse
+from .attention.scoreAttentionPatchedSynapse import ScoreAttentionPatchedSynapse
+
+

ngclearn-3.2.4/ngclearn/components/synapses/attention/__init__.py

@@ -0,0 +1,5 @@
+## attention patched synaptic components
+from .attentionPatchedSynapse import AttentionPatchedSynapse
+from .hebbianAttentionPatchedSynapse import HebbianAttentionPatchedSynapse
+from .scoreAttentionPatchedSynapse import ScoreAttentionPatchedSynapse
+

ngclearn-3.2.4/ngclearn/components/synapses/attention/attentionPatchedSynapse.py

@@ -0,0 +1,295 @@
+# %%
+from jax import random, numpy as jnp, nn
+from ngclearn.components.jaxComponent import JaxComponent
+from ngclearn.utils.distribution_generator import DistributionGenerator
+from ngcsimlib.logger import info
+from ngclearn import compilable, Compartment
+
+def _create_attention_synapses(key, shape, n_patch, use_value=True, init_type=DistributionGenerator.fan_in_gaussian()):
+    key, *subkey = random.split(key, 10)
+
+    patch_shape = (shape[0] // n_patch, shape[1] // n_patch)
+    di, dj = patch_shape
+    n_weight = 3 if use_value else 2        ## synapses: query, key, value (optional)
+
+    weight_shape = (di, n_weight * dj)
+    weights = jnp.zeros(weight_shape)
+    w_masks = jnp.zeros(weight_shape)
+
+    for i in range(n_weight): ## W = (Wq, Wk, Wv)
+        start_j = i * dj
+        end_j = (i + 1) * dj
+
+        weights = weights.at[:, start_j : end_j].set(init_type(patch_shape, subkey[i]))
+        w_masks = w_masks.at[:, start_j : end_j].set(jnp.ones(patch_shape))
+
+    return weights, patch_shape, w_masks
+
+def _split_heads(t, n_heads):
+    """
+    (batch_size, n_patch, d) -> (batch_size, n_heads, n_patch, d / n_heads)
+    """
+    return t.reshape(t.shape[0], t.shape[1], n_heads, -1).transpose(0, 2, 1, 3)
+
+def _merge_heads(t):
+    """
+    (batch_size, n_heads, n_patch, d / n_heads) -> (batch_size, n_patch, d)
+    """
+    return t.transpose(0, 2, 1, 3).reshape(t.shape[0], t.shape[2], -1)
+
+def _attend(inputs, W, b, n_patch, n_heads=1, use_value=True):
+    """
+    (zF W + b) -> y = softmax(q k.T / sqrt(d_h)) v
+    """
+    zF = inputs.reshape(inputs.shape[0], n_patch, -1)  ## (B, n, d)
+    q_k_v = (zF @ W) + b                               ## q_k_v: (B, n, d)
+
+    if use_value:
+        q, k, v = jnp.split(q_k_v, 3, axis=2)          ## q, k v projections of every patch
+    else:
+        q, k = jnp.split(q_k_v, 2, axis=2)             ## v = zF
+        v = zF
+
+    q = _split_heads(q, n_heads)                       ## Reshape: (B, n, d) -> (B, H, n, d/H)
+    k = _split_heads(k, n_heads)                       ## Reshape: (B, n, d) -> (_, H, _, d/H)
+    v = _split_heads(v, n_heads)                       ## Reshape: (B, n, d) -> (_, H, _, d/H)
+    ## attention score
+    rd = jnp.sqrt(q.shape[3])
+    s = (q @ k.transpose(0, 1, 3, 2)) / rd             ## (B, H, n, d/H) @ (B, H, d/H, n)
+    ## probability
+    A = nn.softmax(s, axis=-1)
+    y = _merge_heads(A @ v)                            ## Reshape.T: (B, H, n, d/H) -> (B, n, d)
+    return y.reshape(y.shape[0], -1), (zF, q, k, v, s, A)
+
+def _attend_backward(e, attention_params, n_heads=1):
+    """
+    d_qkv = (chain rule through softmax(q k.T / sqrt(d_h)) v)
+    Passes a post-synaptic signal back through the attention mixing to the (q, k, v) projections of every patch
+    """
+    zF, q, k, v, s, A = attention_params
+    rd = jnp.sqrt(q.shape[3])
+    e = _split_heads(e.reshape(e.shape[0], zF.shape[1], -1), n_heads)   ## Reshape: (B, n, d) -> (_, H, _, d/H)
+    dv = A.transpose(0, 1, 3, 2) @ e
+    dA = e @ v.transpose(0, 1, 3, 2)
+    dS = A * (dA - jnp.sum(dA * A, axis=-1, keepdims=True)) / rd
+    dq = dS @ k
+    dk = dS.transpose(0, 1, 3, 2) @ q
+    return jnp.concatenate([_merge_heads(dq),     ## Reshape.T: (B, H, n, d/H) -> (B, n, d)
+                            _merge_heads(dk),     ## Reshape.T: (B, H, n, d/H) -> (B, n, d)
+                            _merge_heads(dv)],    ## Reshape.T: (B, H, n, d/H) -> (B, n, d)
+                           axis=2)
+
+def _backtransmit(e, attention_params, W, n_heads=1, use_value=True):
+    """
+    pre_out = dq * Wq.T + dk * Wk.T + dv (v = zF)   (chain rule through the weights)
+    """
+    d_qkv = _attend_backward(e, attention_params, n_heads)
+    dzF = d_qkv[:, :, 0:W.shape[1]] @ W.T
+    if not use_value: ## (v = zF): the value signal passes straight back
+        dzF = dzF + d_qkv[:, :, W.shape[1]:]
+    return dzF.reshape(dzF.shape[0], -1)
+
+class AttentionPatchedSynapse(JaxComponent): ## base attention patched synaptic cable
+    """
+    A patched self-attention synaptic cable that lets the patches (tokens) of a layer attend to one another, i.e., it
+    transforms inputs via multi-head self-attention through a single set of query, key, and value synapses that is
+    shared by every patch, no form of synaptic evolution/adaptation is in-built to this component.
+
+    | --- Synapse Compartments: ---
+    | inputs - input (takes in external signals)
+    | outputs - output signals (transformation induced by synapses)
+    | project_input - input of the projection pathway (takes in external signals)
+    | project_output - output signals of the projection pathway (transformation induced by synapses)
+    | post_in - post-synaptic input of the backward pathway (takes in external signals)
+    | pre_out - pre-synaptic output signals of the backward pathway (post_in carried back across the synapses)
+    | weights - current value matrix of synaptic efficacies, the query, key, and value synapses W = (Wq, Wk, Wv)
+    | biases - current value vector of synaptic bias values, b = (bq, bk, bv)
+    | key - JAX PRNG key
+
+    Args:
+        name: the string name of this cell
+
+        shape: tuple specifying shape of this synaptic cable (usually a 2-tuple
+            with number of inputs by number of outputs), every patch gets
+            shape[0] / n_patch input and shape[1] / n_patch output
+            features (these must be equal if use_value=False)
+
+        n_patch: The number of patches (tokens) in each layer (Default: 1)
+
+        n_heads: The number of attention heads, each head attends over
+            (shape[1] / n_patch) / n_heads features (Default: 1)
+
+        use_value: if False, the value synapses are fixed to the identity (v = zF)
+            and are left out of `weights`, i.e., W = (Wq, Wk) (Default: True)
+
+        weight_init: a kernel to drive initialization of this synaptic cable's values;
+            typically a tuple with 1st element as a string calling the name of
+            initialization to use
+
+        bias_init: a kernel to drive initialization of biases for this synaptic cable
+            (Default: None, which turns off/disables biases)
+
+        resist_scale: a fixed scaling factor to apply to synaptic transform
+            (Default: 1.), i.e., yields: out = attention(in; W, b) * Rscale
+
+        p_conn: probability of a connection existing (default: 1.); setting
+            this to < 1. will result in a sparser synaptic structure
+
+        batch_size: the size of each mini batch
+    """
+
+    def __init__(self, name, shape, n_patch=1, n_heads=1, use_value=True, weight_init=None, bias_init=None,
+            resist_scale=1., p_conn=1., batch_size=1, **kwargs
+    ):
+        super().__init__(name, **kwargs)
+
+        self.Rscale = resist_scale
+        self.batch_size = batch_size
+        self.weight_init = weight_init
+        self.bias_init = bias_init
+
+        self.n_patch = n_patch
+        self.n_heads = n_heads
+        self.use_value = use_value
+
+        tmp_key, *subkeys = random.split(self.key.get(), 4)
+        if self.weight_init is None:
+            info(self.name, "is using default weight initializer!")
+            self.weight_init = DistributionGenerator.fan_in_gaussian()
+
+        weights, self.patch_shape, self.w_masks = _create_attention_synapses(
+            key=tmp_key, shape=shape, n_patch=self.n_patch, use_value=self.use_value,
+            init_type=self.weight_init
+        )
+
+        self.shape = shape ## input x output same for all weights (Wq, Wk, Wv)
+
+        if 0. < p_conn < 1.: ## only non-zero and <1 probs allowed
+            mask = random.bernoulli(subkeys[1], p=p_conn, shape=weights.shape)
+            weights = weights * mask ## sparsify matrix
+
+        ## Compartment setup
+        self.weights = Compartment(weights)
+
+        preVals = jnp.zeros((self.batch_size, self.shape[0]))
+        postVals = jnp.zeros((self.batch_size, self.shape[1]))
+
+        self.inputs = Compartment(preVals)
+        self.outputs = Compartment(postVals)
+
+        self.project_input = Compartment(preVals)
+        self.project_output = Compartment(postVals)
+
+        self.post_in = Compartment(postVals)
+        self.pre_out = Compartment(preVals)
+
+        ## Set up (optional) bias values
+        if self.bias_init is None:
+            info(self.name, "is using default bias value of zero (no bias kernel provided)!")
+        self.biases = Compartment(self.bias_init((1, weights.shape[1]), subkeys[2]) if bias_init else 0.0)
+
+    @compilable
+    def advance_state(self):
+        # Get the variables
+        weights = self.weights.get()
+        biases = self.biases.get()
+
+        ################### inputs  →  attention(W)= A @ v  → outputs
+        inputs = self.inputs.get()  ## zF
+        ## Compute (zF @ W) -> softmax(q k.T / sqrt(d_h)) v
+        outputs, attention_params = _attend(inputs, weights, biases, self.n_patch, self.n_heads, self.use_value)
+        _, q, k, v, s, A = attention_params
+        outputs = outputs * self.Rscale
+        ## Update outputs compartment
+        self.outputs.set(outputs)
+
+
+        ###################    post_in →  (d_qkv @ W.T)  → pre_out
+        post_in = self.post_in.get()
+        ## Compute attention(post_in, W.T)
+        pre_out = _backtransmit(post_in, attention_params, weights, self.n_heads, self.use_value)
+        ## Update pre_out compartment
+        self.pre_out.set(pre_out)
+
+
+        ################### project_input →  attention(W)  → project_output
+        project_input = self.project_input.get()
+        ## Compute attention(project_input, W)
+        project_output, _ = _attend(project_input, weights, biases, self.n_patch, self.n_heads, self.use_value)
+        project_output = project_output * self.Rscale
+        ## Update project_output compartment
+        self.project_output.set(project_output)
+
+    @compilable
+    def reset(self): ## closed, no-batch argument reset
+        ## write reset command to call inner batched_reset command
+        self.batched_reset(batch_size=self.batch_size) ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size): ## open, batch argument reset
+        preVals = jnp.zeros((batch_size, self.shape[0]))
+        postVals = jnp.zeros((batch_size, self.shape[1]))
+
+        # BUG: the self.inputs here does not have the targeted field
+        # NOTE: Quick workaround is to check if targeted is in the input or not
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals)
+        self.outputs.set(postVals)
+        self.post_in.set(postVals)
+        self.pre_out.set(preVals)
+        self.project_input.set(preVals)
+        self.project_output.set(postVals)
+
+    @classmethod
+    def help(cls): ## component help function
+        properties = {
+            "synapse_type": "AttentionPatchedSynapse - performs a synaptic transformation "
+                            "of inputs to produce output signals (e.g., a multi-head "
+                            "self-attention transformation across the patches of a layer)"
+        }
+        compartment_props = {
+            "inputs":
+                {"inputs": "Takes in external input signal values",
+                 "project_input": "Takes in external input signal values",
+                 "post_in": "Takes in external input signal values",
+                 },
+            "states":
+                {"weights": "Synapse efficacy/strength parameter values (query, key, and value synapses)",
+                 "biases": "Base-rate/bias parameter values (query, key, and value biases)",
+                 "key": "JAX PRNG key"},
+            "outputs":
+                {"outputs": "Output of synaptic transformation",
+                 "project_output": "Output of synaptic transformation",
+                 "pre_out": "Output of synaptic transformation",
+                 },
+        }
+        hyperparams = {
+            "shape": "Overall shape of synaptic cable, number inputs x number outputs",
+            "n_patch": "The number of patches (tokens) in each layer",
+            "n_heads": "The number of attention heads",
+            "use_value": "Are the value synapses adaptive? (otherwise, they are fixed to the identity)",
+            "batch_size": "Batch size dimension of this component",
+            "weight_init": "Initialization conditions for synaptic weight (W) values",
+            "bias_init": "Initialization conditions for bias/base-rate (b) values",
+            "resist_scale": "Resistance level scaling factor (Rscale), applied to output of transformation",
+            "w_masks": "weight mask matrix",
+            "p_conn": "Probability of a connection existing (otherwise, it is masked to zero)"
+        }
+        info = {cls.__name__: properties,
+                "compartments": compartment_props,
+                "dynamics": "outputs = [softmax((inputs * Wq + bq) * (inputs * Wk + bk).T / sqrt(d)) * (inputs * Wv + bv)] * Rscale",
+                "hyperparameters": hyperparams}
+        return info
+
+if __name__ == '__main__':
+    from ngcsimlib.context import Context
+
+    with Context("Bar") as bar:
+        Wab = AttentionPatchedSynapse("Wab", (9, 9), 3)
+    print(Wab)
+
+
+
+
+
+
+