diff --git a/README.md b/README.md
index 731a02b8..f7312883 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ ngc-learn requires:
1) Python (>=3.10)
2) NumPy (>=1.22.0)
3) SciPy (>=1.7.0)
-4) ngcsimlib (>=3.1.0), (visit official page here)
+4) ngcsimlib (>=3.1.1), (visit official page here)
5) JAX (>=0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
---
-ngc-learn 3.2.0 and later require Python 3.10 or newer as well as ngcsimlib >=3.1.0.
+ngc-learn 3.2.1 and later require Python 3.10 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).
@@ -72,7 +72,7 @@ Python 3.11.4 (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.0'
+'3.2.1'
```
Note: For access to the previous Tensorflow-2 version of ngc-learn (of
@@ -119,7 +119,7 @@ $ python install -e .
**Version:**
-3.2.0
+3.2.1
Author:
Alexander G. Ororbia II
diff --git a/ngclearn/components/__init__.py b/ngclearn/components/__init__.py
index 4e32c6c4..2a622205 100644
--- a/ngclearn/components/__init__.py
+++ b/ngclearn/components/__init__.py
@@ -50,6 +50,7 @@
from .synapses.hebbian.eventSTDPSynapse import EventSTDPSynapse
from .synapses.hebbian.BCMSynapse import BCMSynapse
from .synapses.hebbian.inhibitorySTDPSynapse import InhibitorySTDPSynapse
+from .synapses.hebbian.ojaTensorSynapse import OjaTensorSynapse
from .synapses.competitive.SOMSynapse import SOMSynapse
from .synapses.STPDenseSynapse import STPDenseSynapse
from .synapses.exponentialSynapse import ExponentialSynapse
diff --git a/ngclearn/components/synapses/__init__.py b/ngclearn/components/synapses/__init__.py
index 5bb09fbd..2f7b829b 100644
--- a/ngclearn/components/synapses/__init__.py
+++ b/ngclearn/components/synapses/__init__.py
@@ -17,6 +17,7 @@
from .hebbian.eventSTDPSynapse import EventSTDPSynapse
from .hebbian.BCMSynapse import BCMSynapse
from .hebbian.inhibitorySTDPSynapse import InhibitorySTDPSynapse
+from .hebbian.ojaTensorSynapse import OjaTensorSynapse
from .mpsSynapse import MPSSynapse
### dense competitive synaptic components/elements
from .competitive.SOMSynapse import SOMSynapse
diff --git a/ngclearn/components/synapses/hebbian/__init__.py b/ngclearn/components/synapses/hebbian/__init__.py
index 1a27c09a..9f69e820 100644
--- a/ngclearn/components/synapses/hebbian/__init__.py
+++ b/ngclearn/components/synapses/hebbian/__init__.py
@@ -6,4 +6,5 @@
from .BCMSynapse import BCMSynapse
from .gerstnerHebbianSynapse import GerstnerHebbianSynapse ## Taylor-expansion Hebbian model
from .inhibitorySTDPSynapse import InhibitorySTDPSynapse ## iSTDP rule (trace-based)
+from .ojaTensorSynapse import OjaTensorSynapse ## generalized Oja's synaptic cable
diff --git a/ngclearn/components/synapses/hebbian/ojaTensorSynapse.py b/ngclearn/components/synapses/hebbian/ojaTensorSynapse.py
new file mode 100644
index 00000000..3ee327fb
--- /dev/null
+++ b/ngclearn/components/synapses/hebbian/ojaTensorSynapse.py
@@ -0,0 +1,229 @@
+from jax import random, numpy as jnp, jit
+#from functools import partial
+from ngclearn import compilable
+from ngclearn import Compartment
+from ngclearn.components.synapses import SparseTensorSynapse
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
+
+class OjaTensorSynapse(SparseTensorSynapse):
+ """
+ A sparse tensor-synaptic cable that is adapted via Hebbian or Oja-plasticity. Note this component cable implements
+ a full, locally-connected structure or an unshared convolutional synaptic tensor structural component (cable).
+
+ | --- Synapse Input Compartments: ---
+ | inputs - input (takes in external signals)
+ | --- Synapse State Compartments: ---
+ | weights - current value matrix of synaptic efficacies (i.e., the strength values)
+ | biases - current value vector of synaptic bias values
+ | --- Synapse Output Compartments: ---
+ | outputs - output signals
+
+ | References:
+ | Oja, E., 1982. Simplified neuron model as a principal component analyzer. Journal of mathematical biology,
+ | 15(3), pp.267-273.
+
+ Args:
+ name: the string name of this cell
+
+ n_in_streams: total number of incoming streams
+
+ K_local: input local block feature size
+
+ O_local: output local block feature size (note: total number of output streams is computed internally)
+
+ P_l: local window size (how many input streams to grab), i.e., maximum connection window size
+ (defines tensor shape) (Default: 1)
+
+ stride: stride factor for locally-connected structure to "skip over" (in terms of incoming streams);
+ stride step between successive output receptive field streams (Default: 1)
+
+ convergent_factor: if > 0, windows contract symmetrically towards center of layer
+ (simulates foveation or variable resolution) (Default: 0)
+
+ dilation: dilation factor; gap between input indices inside a block (1 = contiguous) (Default: 1)
+
+ invert_conn: if True, this tensor shape will internally "transpose" itself to formulate the appropriate
+ locally-connected inverted/transposed structure (note that this means the shape will no longer be what
+ the constructor's argument dictate - it will be the effective transpose of these arguments) (Default: False)
+
+ eta: global learning rate
+
+ oja_factor: if > 0, this triggers the Oja-correction factor to Hebbian plasticity (Default: 0)
+
+ sign_value: multiplicative factor to apply to final synaptic update before it is applied to synapses; this is
+ useful if gradient descent style optimization is required (as Hebbian rules typically yield adjustments
+ for ascent)
+
+ normalize: if True, this synaptic tensor will normalize its internal blocks (Default: False)
+
+ norm_axis: axis upon which block norms are computed (Default: 1)
+
+ 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)
+
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * in) * g_conduct_factor) + bias
+
+ p_conn: probability of a connection existing (default: 1.); setting this to < 1 and > 0. will result in a
+ sparser synaptic structure (lower values yield sparse structure)
+
+ use_block_matrix_format: if True, this tensor synapse resorts to a memory-intensive block-matrix
+ computational format (Default: False)
+
+ """
+
+ def __init__(
+ self,
+ name,
+ n_in_streams, ## total incoming streams (nlm1_streams)
+ K_local, ## input local block feature size
+ O_local, ## output local block feature size
+ P_l=1, ## local window size (how many input streams to grab)
+ stride=1,
+ convergent_factor=0.,
+ dilation=1,
+ invert_conn=False,
+ eta=0.,
+ oja_factor=1., ## if <= 0, will disable Oja's correction term
+ sign_value=1.,
+ normalize=False,
+ norm_axis=1,
+ weight_init=None,
+ bias_init=None,
+ p_conn=1.,
+ resist_scale=1.,
+ batch_size=1,
+ use_block_matrix_format=False, ## if True, triggers block-matrix implementation
+ **kwargs
+ ):
+ super().__init__(
+ name,
+ n_in_streams=n_in_streams, # total incoming streams (nlm1_streams)
+ K_local=K_local, ## input local block feature size
+ O_local=O_local, ## output local block feature size
+ P_l=P_l, ## local window size (how many input streams to grab?)
+ stride=stride,
+ convergent_factor=convergent_factor,
+ dilation=dilation,
+ invert_conn=invert_conn,
+ normalize=normalize,
+ norm_axis=norm_axis,
+ use_block_matrix_format=use_block_matrix_format,
+ weight_init=weight_init,
+ bias_init=bias_init,
+ resist_scale=resist_scale,
+ p_conn=p_conn,
+ batch_size=batch_size,
+ **kwargs
+ )
+
+ ## store cable framework properties
+ self.eta = eta
+ self.sign_value = sign_value
+ self.oja_factor = oja_factor
+
+ # if self.normalize:
+ # P, S, Klocal, Olocal = self.shape
+ # norm_axis = 1 #0 #1
+ # Weffdim = Klocal * P
+ # if norm_axis == 1:
+ # Weffdim = Olocal * self.n_out_streams #Weffdim = Olocal
+ # weight_init = dist.gaussian(mean=0., std=float(1.0 / jnp.sqrt(Weffdim)))
+ # weights = weight_init(self.shape, self.key.get())
+ # self.weights.set(weights)
+
+ preVals = jnp.zeros((self.batch_size, self.io_shape[0]))
+ postVals = jnp.zeros((self.batch_size, self.io_shape[1]))
+ self.pre = Compartment(preVals)
+ self.post = Compartment(postVals)
+ self.dWeights = Compartment(self.weights.get() * 0)
+ #self.dBiases = Compartment(jnp.zeros(shape[1]))
+ self.variance_term = Compartment(postVals)
+
+ @compilable
+ def evolve(self, dt):
+ ## get compartment variables
+ pre = self.pre.get()
+ post = self.post.get()
+ weights = self.weights.get()
+ # biases = self.biases.get()
+ key, *subkeys = random.split(self.key.get(), 3) ## generate keys for noise samples
+ self.key.set(key)
+
+ B = pre.shape[0]
+ if not self.use_block_matrix_format: ## use tensor format
+ conn_map = self.connectivity_map ## get connectivity structure
+ inputs = pre
+ outputs = post #+ var_term
+ P, S, K_local, O_local = weights.shape
+
+ ## 1st term - calculate (Hebbian) correlation term
+ ## gather inputs and transpose to: (P, B, S, K_local)
+ patched_inputs = inputs.reshape(B, -1, K_local)
+ pre = jnp.transpose(patched_inputs[:, conn_map, :], (2, 0, 1, 3))
+ ## reshape outputs to structure: (B, S, O_local)
+ post = outputs.reshape(B, S, O_local)
+ ## use None axes to broadcast (P, B, S, K_local, 1) * (1, B, S, 1, O_local)
+ ### resulting shape: (P, B, S, K_local, O_local)
+ outer_product_grid = pre[..., None] * post[None, :, :, None, :]
+ ## sum over the batch dimension (axis 1) -> (P, S, K_local, O_local)
+ dWeights_hebb = jnp.sum(outer_product_grid, axis=1)
+
+ ## 2nd term - apply (optional) Oja's correction factor
+ dWeights_oja = 0.
+ if self.oja_factor > 0.:
+ ## compute squared postsynaptic activation; (B, S, O_local)
+ post_squared = jnp.square(post)
+ ## average across batch dimension to get expected decay weight per neuron; (S, O_local)
+ mean_post_squared = jnp.sum(post_squared, axis=0)
+ ## broadcast decay against current synaptic weights
+ ### (1, S, 1, O_local) * (P, S, K_local, O_local) -> (P, S, K_local, O_local)
+ dWeights_oja = -(mean_post_squared[None, :, None, :] * weights)
+
+ dWeights = dWeights_hebb + dWeights_oja
+ else: ## use block-matrix format (memory-intensive)
+ dWeights_hebb = pre.T @ post ## classical Hebbian term
+ dWeights_oja = 0.
+ if self.oja_factor > 0.: ### 2nd, apply Oja's (corrective) factor
+ dWeights_oja = -(post[:, None] ** 2 * weights) ## Oja-correction / stability term
+ dWeights_oja = jnp.sum(dWeights_oja, axis=0) ## calculate full (tensor) Oja's correction
+ dWeights = (dWeights_hebb + dWeights_oja) * (weights != 0.) ## calculate full (tensor) update
+
+ weights = weights + dWeights * self.sign_value
+
+ ## apply post-projection / normalization steps
+ if self.norm_trigger.get() > 0.:
+ if not self.use_block_matrix_format:
+ weights = SparseTensorSynapse._normalize_outgoing_weights(weights)
+ else:
+ weights = SparseTensorSynapse._normalize_block_matrix_outgoing_weights(
+ weights,
+ self.connectivity_map,
+ self.n_in_streams,
+ self.K_local,
+ self.O_local,
+ self.n_out_streams,
+ axis=1, #0, # 1
+ order=2,
+ norm_targ=1.
+ )
+
+ ## update compartments
+ self.weights.set(weights)
+ self.dWeights.set(dWeights)
+
+ @compilable
+ def reset(self): ## reset core components/statistics
+ if not self.inputs.targeted:
+ self.inputs.set(self.inputs.get() * 0)
+ self.outputs.set(self.outputs.get() * 0) # outputs
+ self.pre.set(self.pre.get() * 0) # pre
+ self.post.set(self.post.get() * 0) # post
+ self.dWeights.set(self.dWeights.get() * 0) # dW
+ #self.dBiases.set(self.dBiases.get() * 0) # db
+ self.variance_term.set(self.variance_term.get() * 0)
+
diff --git a/ngclearn/components/synapses/sparseTensorSynapse.py b/ngclearn/components/synapses/sparseTensorSynapse.py
index aea25558..d1a85a38 100644
--- a/ngclearn/components/synapses/sparseTensorSynapse.py
+++ b/ngclearn/components/synapses/sparseTensorSynapse.py
@@ -12,9 +12,9 @@
########################################################################################################################
## helper-functions for sparse-synaptic tensors
def _make_connectivity_map(
- P,
- S,
- stride=1,
+ P, ## size of sliding input connectivity window
+ S, ## number of output streams/maps/groups
+ stride=1, ## stream skip (Default: 1)
padding=0,
dilation=1,
convergent_factor=0.0,
@@ -58,7 +58,10 @@ def _make_connectivity_map(
center = S / 2.0
distance_from_center = abs(s - center) / center
## shrink window proportionally, ensuring it uses at least 1 connection
- local_P = max(1, int(P * (1.0 - convergent_factor * distance_from_center)))
+ #local_P = max(1, int(P * (1.0 - convergent_factor * distance_from_center)))
+ local_P = max(
+ 1, round(P * (1.0 - convergent_factor * distance_from_center))
+ )
else:
local_P = P
@@ -78,13 +81,6 @@ def _make_connectivity_map(
conn_map[s, p] = rng.integers(0, total_inputs)
return jnp.array(conn_map)
-# def make_sliding_connectivity_map(P, S, stride=1, padding=0):
-# conn_map = np.zeros((S, P), dtype=np.int32)
-# for s in range(S):
-# start_idx = s * stride - padding
-# conn_map[s, :] = np.arange(start_idx, start_idx + P)
-# return jnp.array(conn_map)
-
@partial(jit, static_argnums=(2, 3, 4, 5))
def _reconstruct_global_2d_matrix(
weights: jnp.ndarray,
@@ -138,6 +134,70 @@ def _reconstruct_global_2d_matrix(
].set(weights.ravel())
return global_2d_matrix
+
+def precompute_transpose_indices(
+ forward_conn_map,
+ backward_conn_map,
+ forward_shape,
+ backward_shape
+):
+ """
+ Pre-computes a flat index gather array mapping a forward unshared tensor to its exact transposed
+ backward tensor configuration.
+
+ Forward Shape: (P_f, S_f, K_f, O_f)
+ Backward Shape: (P_b, S_b, K_b, O_b) -> Typically K_b == O_f and O_b == K_f
+ """
+ P_f, S_f, K_f, O_f = forward_shape
+ P_b, S_b, K_b, O_b = backward_shape
+
+ ## create empty template of destination tensor layout
+ ### we populate this with absolute memory addresses of forward tensor
+ backward_flat_indices = np.full((P_b, S_b, K_b, O_b), -1, dtype=np.int32)
+ ## loop over every single target slot in backward transpose tensor
+ for p_b in range(P_b):
+ for s_b in range(S_b):
+ ## look up which input stream this backward slot is "listening" to
+ input_stream = backward_conn_map[s_b, p_b]
+ if input_stream == -1:
+ continue ## skip dead padding zones
+ ## in forward pass, this destination 's_b' was an input stream,
+ ### and input_stream was destination output stream ('s_f')
+ s_f = int(input_stream)
+ ## scan forward connectivity map for this stream to locate
+ ### where it connects back to backward stream source 's_b'
+ forward_inputs = forward_conn_map[s_f, :]
+ matches = np.where(forward_inputs == s_b)[0]
+
+ if len(matches) > 0:
+ p_f = int(matches[0])
+
+ ## slices connect across local neural blocks:
+ ### forward K_f maps to backward O_b; forward O_f maps to backward K_b
+ for k_b in range(K_b):
+ for o_b in range(O_b):
+ ## map back to corresponding forward coordinates
+ k_f = o_b
+ o_f = k_b
+ ## compute absolute 1D flat memory index inside forward tensor
+ forward_flat_idx = (
+ p_f * (S_f * K_f * O_f) +
+ s_f * (K_f * O_f) +
+ k_f * O_f +
+ o_f
+ )
+ ## store coordinate memory location in lookup block
+ backward_flat_indices[p_b, s_b, k_b, o_b] = forward_flat_idx
+
+ ## flatten the lookup block into a 1D index array for quick GPU gathering
+ flat_gather_indices = backward_flat_indices.flatten()
+ ## create a mask array to cleanly zero-out entries corresponding to dead padding zones (-1)
+ valid_mask = (flat_gather_indices != -1).astype(np.float32)
+ ## clip indices to 0 so jnp.take doesn't throw out-of-bounds tracking faults on padding slots
+ flat_gather_indices = np.clip(flat_gather_indices, 0, None)
+
+ return jnp.array(flat_gather_indices), jnp.array(valid_mask)
+
def _make_backwards_connectivity_map(
forward_conn_map, total_input_streams
):
@@ -257,32 +317,41 @@ def __init__(
## set up internal structure/topology
### run geometry math internally first; if invert_conn is True, 'incoming' streams to this component are actually
### errors from upper layer, meaning 'in_streams' is upper layer's output streams (S_forward)
+ ## standard forward-projection synaptic cable
+ dims = SparseTensorSynapse.initialize_layer_geometry(
+ n_in_streams,
+ K_local,
+ O_local,
+ P_l,
+ stride,
+ auto_pad=False,
+ dilation_l=dilation,
+ convergent_factor_l=convergent_factor
+ )
+
if invert_conn:
- dims = SparseTensorSynapse.initialize_layer_geometry(
- n_in_streams,
- K_local,
- O_local,
- P_l,
- stride,
- auto_pad=False,
- dilation_l=dilation,
- convergent_factor_l=convergent_factor
+ # 1. Build the TRUE forward connectivity map matching your architecture's dimensions
+ forward_connectivity_map = _make_connectivity_map(
+ P_l, dims["total_num_output_streams"], stride, convergent_factor=convergent_factor, dilation=dilation
)
+ _old_S = dims["total_num_output_streams"]
+ _old_forward_shape = (P_l, _old_S, K_local, O_local)
total_forward_inputs = n_in_streams
total_forward_outputs = dims["total_num_output_streams"]
- ## build base forward connectivity map so this can later be "inverted"
- f_map = _make_connectivity_map(
- P_l, S=n_in_streams, stride=stride, convergent_factor=convergent_factor, dilation=dilation
+ # 2. CRITICAL FIX: Invert the EXACT SAME map we just generated above!
+ # Do NOT call _make_connectivity_map a second time with different stream sizes.
+ backward_connectivity_map, P_back = _make_backwards_connectivity_map( ## invert map
+ forward_connectivity_map, total_forward_inputs
)
- self.connectivity_map, P_back = _make_backwards_connectivity_map(f_map, total_forward_inputs) ## invert map
+ self.connectivity_map = backward_connectivity_map
## explicitly compile exact shapes that ngc-learn expects
- self.n_in_streams = total_forward_outputs ## e.g., 1 (number of incoming error maps)
- self.n_out_streams = total_forward_inputs ## e.g., 3 (number of feedback target maps)
- self.K_local = O_local # Error size (9)
- self.O_local = K_local # Hidden feature target size (16)
+ self.n_in_streams = total_forward_outputs
+ self.n_out_streams = total_forward_inputs
+ self.K_local = O_local
+ self.O_local = K_local
# Final 4D tensor shape for backwards weights
calculated_shape = (P_back, total_forward_inputs, O_local, K_local)
@@ -291,19 +360,16 @@ def __init__(
# Clear window strides for pure backward routing execution
self.stride = 0
self.dilation = 1
- else:
- ## standard forward-projection synaptic cable
- dims = SparseTensorSynapse.initialize_layer_geometry(
- n_in_streams,
- K_local,
- O_local,
- P_l,
- stride,
- auto_pad=False,
- dilation_l=dilation,
- convergent_factor_l=convergent_factor
- )
+ # 3. Both maps are now in perfect structural harmony
+ self.transpose_indices, self.transpose_mask = precompute_transpose_indices(
+ forward_connectivity_map,
+ backward_connectivity_map,
+ _old_forward_shape,
+ calculated_shape
+ )
+ else:
+ self.transpose_indices = self.transpose_mask = None
self.n_in_streams = n_in_streams # dims["total_input_streams"] # nlm1_streams
self.n_out_streams = dims["total_num_output_streams"] # Computed S
self.K_local = K_local
@@ -341,14 +407,14 @@ def __init__(
self.norm_trigger = Compartment(jnp.zeros((1, 1)))
if self.normalize:
self.norm_trigger.set(self.norm_trigger.get() + 1)
- P, S, Klocal, Olocal = self.shape
- norm_axis = 1 #0 #1
- Weffdim = Klocal * P
- if norm_axis == 1:
- Weffdim = Olocal * self.n_out_streams #Weffdim = Olocal
- weight_init = dist.gaussian(mean=0., std=float(1.0 / jnp.sqrt(Weffdim)))
- weights = weight_init(self.shape, self.key.get())
- self.weights.set(weights)
+ # P, S, Klocal, Olocal = self.shape
+ # norm_axis = 1 #0 #1
+ # Weffdim = Klocal * P
+ # if norm_axis == 1:
+ # Weffdim = Olocal * self.n_out_streams #Weffdim = Olocal
+ # weight_init = dist.gaussian(mean=0., std=float(1.0 / jnp.sqrt(Weffdim)))
+ # weights = weight_init(self.shape, self.key.get())
+ # self.weights.set(weights)
if self.use_block_matrix_format: ## set up block-diagonal algorithm backend if flagged
block_weights = _reconstruct_global_2d_matrix(
@@ -378,12 +444,13 @@ def __init__(
)
self.weights.set(weights)
-
+ ## set up core compartments
self.initial_weights = Compartment(self.weights.get()) ## store synaptic initial conditions (never updated)
preVals = jnp.zeros((self.batch_size, self.io_shape[0]))
postVals = jnp.zeros((self.batch_size, self.io_shape[1]))
self.inputs.set(preVals)
self.outputs.set(postVals)
+ self.weight_transpose_target = Compartment(self.weights.get())
@staticmethod
def initialize_layer_geometry( ## internal co-routine for tensor-synaptic projection shaping
@@ -470,6 +537,7 @@ def advance_state(self):
biases = self.biases.get()
inputs = self.inputs.get() ## get inputs
conn_map = self.connectivity_map ## get connectivity structure
+ #print(self.name, " ", weights.shape, " ", inputs.shape)
B = inputs.shape[0] ## get batch size for subsequent averaging
if not self.use_block_matrix_format:
@@ -509,7 +577,7 @@ def advance_state(self):
outputs = outputs_3d.reshape(B, -1) ## flatten to 2D formatted output
else:
outputs = inputs @ weights ## block-diagonal/block-matrix multiply
- outputs + biases
+ outputs = (outputs * self.g_conduct_factor) + biases
self.outputs.set(outputs)
@compilable
@@ -530,12 +598,16 @@ def _normalize_outgoing_weights( ## internal tensor normalizer co-routine
### - axis 0 (P) and Axis 2 (K_local) isolate an individual input neuron.
### - axis 1 (S) and Axis 3 (O_local) represent all the places its synapses land.
- if axis == 0:
- ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
- reduce_axes = (0, 2)
- elif axis == 1:
- ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
- reduce_axes = (1, 3)
+ # if axis == 0:
+ # ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
+ # reduce_axes = (0, 2)
+ # elif axis == 1:
+ # ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
+ # reduce_axes = (1, 3)
+ if axis == 1:
+ reduce_axes = 3 ## get col-axis for normalization
+ elif axis == 0:
+ reduce_axes = 2 ## assume row-axis for normalization
else:
raise ValueError("Norm.axis must be 0 or 1.")
@@ -597,12 +669,16 @@ def _normalize_block_matrix_outgoing_weights( ## internal block-matrix normalize
gathered_weights = global_matrix[global_y_indices, global_x_indices]
## compute norms on the isolated valid blocks based on the desired axis
- if axis == 0:
- ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
- reduce_axes = (0, 2)
- elif axis == 1:
- ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
- reduce_axes = (1, 3)
+ # if axis == 0:
+ # ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
+ # reduce_axes = (0, 2)
+ # elif axis == 1:
+ # ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
+ # reduce_axes = (1, 3)
+ if axis == 1:
+ reduce_axes = 3 ## get col-axis for normalization
+ elif axis == 0:
+ reduce_axes = 2 ## assume row-axis for normalization
else:
raise ValueError("Norm.axis must be 0 or 1.")
@@ -622,3 +698,38 @@ def _normalize_block_matrix_outgoing_weights( ## internal block-matrix normalize
].set(normalized_gathered.ravel())
return new_global_matrix
+ @compilable
+ def tie_transpose_weights(self):
+ current_weights = self.weights.get()
+ weights_to_transpose_absorb = self.weight_transpose_target.get()
+
+ transposed_tensor = SparseTensorSynapse._transpose_synaptic_tensor(
+ weights_to_transpose_absorb,
+ self.transpose_indices,
+ self.transpose_mask,
+ current_weights.shape
+ )
+ self.weights.set(transposed_tensor)
+
+ @staticmethod
+ def _transpose_synaptic_tensor( ## internal lc-tensor transpose co-routine
+ forward_weights,
+ flat_gather_indices,
+ valid_mask,
+ backward_shape
+ ):
+ # Executes a 4D locally-connected tensor transpose operation instantly.
+ # Arguments:
+ # forward_weights_flat: The RAW forward tensor, flattened to 1D via .flatten()
+ # flat_gather_indices: The pre-computed index lookup array from Stage 1
+ # valid_mask: The pre-computed padding gate array from Stage 1
+ # backward_shape: Static tuple indicating target shape (P_b, S_b, K_b, O_b)
+ forward_weights_flat = forward_weights.flatten()
+ ## memory address gathering: pull elements into correct positions instantly
+ transposed_flat = jnp.take(forward_weights_flat, flat_gather_indices)
+ ## masking: enforce absolute zero over dead-zone padding slots
+ transposed_flat = transposed_flat * valid_mask
+ ## reshape: snap 1D gathered block back into true 4D backward tensor dimensions
+ return jnp.reshape(transposed_flat, backward_shape)
+
+
diff --git a/ngclearn/utils/analysis/effective_dim.py b/ngclearn/utils/analysis/effective_dim.py
index d4327f4e..c427c47e 100644
--- a/ngclearn/utils/analysis/effective_dim.py
+++ b/ngclearn/utils/analysis/effective_dim.py
@@ -1,4 +1,5 @@
from functools import partial
+import jax
from jax import numpy as jnp, jit
@partial(jit, static_argnums=[1])
@@ -6,7 +7,8 @@ def participation_ratio(
latent_codes, use_NaN_fallback=False
):
"""
- Calculates the participation ratio coefficient for a set of latent codes
+ Calculates the participation ratio coefficient (also known as the Gini effective
+ dimension) for a set of latent codes.
Args:
latent_codes: a set of (N x D) latent code vectors (one row per vector code)
@@ -23,8 +25,9 @@ def participation_ratio(
tr = jnp.trace(cov)
tr2_cov = tr * tr
- cov2_tr = jnp.trace(cov @ cov)
-
+ ## calc frob-norm squared; NOTE: faster line replaced older un-commented one to right
+ cov2_tr = jnp.sum(jnp.square(cov)) #cov2_tr = jnp.trace(cov @ cov)
+
## this algorithm supports one of two fallback cases
if not use_NaN_fallback: ## use fallback-to-1 eff-dim check
## use JAX-friendly conditional / direct switch to fallback to 1.0.
@@ -33,8 +36,6 @@ def participation_ratio(
##else, use ML-oriented NaN return value fallback
return tr2_cov / cov2_tr if cov2_tr > 0 else float("nan")
-
-
@partial(jit, static_argnums=[1])
def rankme(latent_codes, eps=1e-7):
"""
@@ -56,10 +57,48 @@ def rankme(latent_codes, eps=1e-7):
singular_values = jnp.linalg.svd(latent_codes, compute_uv=False) ## singular values of latent_codes
sum_singular_values = jnp.sum(singular_values) ## L1
sum_S_vals = jnp.where(sum_singular_values > 0.0, sum_singular_values, 1.0)
- p = singular_values / sum_S_vals + eps ## L1-normalized singular value
- shannon_entropy = -jnp.sum(p * jnp.log(p)) ## calc Shannon entropy
+ p = singular_values / (sum_S_vals + eps) ## L1-normalized singular value
+ safe_p = jnp.where(p > 0.0, p, 1.0)
+ shannon_entropy = -jnp.sum(p * jnp.log(safe_p)) ## calc Shannon entropy
+
+ ## compute final exp(Shannon entropy) = effective rank
+ #rankme_score = jnp.exp( ## compute final exp(Shannon entropy) = effective rank
+ # jnp.where(sum_singular_values > 0.0, shannon_entropy, jnp.nan)
+ #)
+ rankme_score = jnp.where(sum_singular_values > 0.0, jnp.exp(shannon_entropy), 1.0)
+ return rankme_score
+
+@partial(jit, static_argnums=[1])
+def stable_rank(latent_codes, num_iters=10): ## power-iterator method
+ """
+ Computes the stable rank via the power iteration method in order to find the
+ top singular value.
+
+ Args:
+ latent_codes: a set of (N x D) latent code vectors (one row per vector code)
+
+ num_iters: number of iterations to run power iterator calculation (Default: 10)
+
+ Returns:
+ scalar measurement of the stable rank (proxy for effective dimension)
+ """
+ Z = latent_codes
+ Zc = Z - Z.mean(axis=0, keepdims=True) ## center codes
+ frobenius_norm_sq = jnp.sum(jnp.square(Zc)) ## calc squared frob-norm
+
+ ## use power iteration to find dominant eigenvector of Zc.T @ Zc
+ ### start w/ random vector:
+ key = jax.random.PRNGKey(0)
+ v = jax.random.normal(key, (Zc.shape[1], 1))
+ v = v / jnp.linalg.norm(v)
+ ## apply standard power iteration loop
+ for _ in range(num_iters):
+ ## v = (Zc.T @ (Zc @ v))
+ v = Zc.T @ (Zc @ v)
+ v = v / jnp.linalg.norm(v)
+ ## compute largest singular value squared (i.e., the Rayleigh quotient):
+ ### sigma_max^2 = ||Zc @ v||^2
+ sigma_max_sq = jnp.sum(jnp.square(Zc @ v))
+ return jnp.where(sigma_max_sq > 0.0, frobenius_norm_sq / sigma_max_sq, 1.0) # stable-rank score
- return jnp.exp(jnp.where(sum_singular_values > 0.0,
- shannon_entropy, ## compute final exp(Shannon entropy) = effective rank
- jnp.nan))
diff --git a/ngclearn/utils/metric_utils.py b/ngclearn/utils/metric_utils.py
index 409a2f44..9de7437f 100755
--- a/ngclearn/utils/metric_utils.py
+++ b/ngclearn/utils/metric_utils.py
@@ -81,11 +81,6 @@ def measure_workload_entropy(spike_count_matrix):
probs = global_neuron_activity / total_epoch_spikes
full_entropy = -jnp.sum(jnp.where(probs > 0, probs * jnp.log2(probs), 0.0)) ## prevents log2(0)
entropy = full_entropy * _flag
- # if total_epoch_spikes > 0:
- # probs = global_neuron_activity / total_epoch_spikes
- # entropy = -jnp.sum(jnp.where(probs > 0, probs * jnp.log2(probs), 0.0)) ## prevents log2(0)
- # else:
- # entropy = 0.0
return (min_act, max_act, mean_act), entropy, dead_neurons
@partial(jit, static_argnums=[1])
@@ -233,8 +228,7 @@ def measure_sparsity(codes, tolerance=0., preserve_batch=True, flip_measure=Fals
| rho_mean(X) = 1/N Sum^N_{i=1} rho(x_i)
- where lower/closer to 0 means codes more sparse and closer to 1 means
- codes are more dense.
+ where lower/closer to 0 means codes more sparse and closer to 1 means codes are denser.
Note that this definition of sparsity aligns with Foldiak's definition of
the ratio of active neurons to inactive ones (assuming binary coding):
diff --git a/ngclearn/utils/viz/synapse_plot.py b/ngclearn/utils/viz/synapse_plot.py
index cae28e5c..b60bc71c 100644
--- a/ngclearn/utils/viz/synapse_plot.py
+++ b/ngclearn/utils/viz/synapse_plot.py
@@ -84,8 +84,8 @@ def visualize_macro_grid( ## more complex filter visualization co-routine
plt.figure(figsize=(10, 10), dpi=300)
## use vmin=-1.0 and vmax=1.0 so that the 1.0 canvas background registers as absolute white
- max_val = 1.
- min_val = -1.
+ max_val = None #1.
+ min_val = None #-1.
if contrast_by_data:
max_val = float(jnp.max(jnp.abs(master_canvas)))
min_val = float(jnp.min(jnp.abs(master_canvas)))
diff --git a/pyproject.toml b/pyproject.toml
index 1c14055a..ecd1b495 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" # using setuptool building engine
[project]
name = "ngclearn"
-version = "3.2.0"
+version = "3.2.1"
description = "Simulation software for building and analyzing computational neuroscience models, brain-inspired computing systems, and NeuroAI agents."
authors = [
{name = "Alexander Ororbia", email = "ago@cs.rit.edu"},
diff --git a/requirements.txt b/requirements.txt
index d68c12aa..ede6cb3e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,10 +2,9 @@ numpy>=1.26.4
scikit-learn>=1.6.1
scipy>=1.14.1
matplotlib>=3.9.4
-# patchify ## note: patchify has issues with pip installation
jax>=0.4.28
jaxlib>=0.4.28
-ngcsimlib>=3.1.0
+ngcsimlib>=3.1.1
imageio>=2.37.0
pandas>=2.2.3
typing_extensions>=4.15.0