This example imports a recurrent SNN (Braille classification) from NIR with Spyx's run-and-return from_nir. The core import + inference works end-to-end (from_nir(ng, x, dt=1e-4) classifies the bundled ds_test.pt). The later cells that rebuild single-layer sub-graphs to compare against snnTorch use manual NIR graph surgery that the current nir's stricter type inference is pickier about; they are kept as reference.

In [1]:

import jax.numpy as jnp
import matplotlib.pyplot as plt
import nir
import numpy as np

# for loading dataset
import torch

import spyx

import jax.numpy as jnp import matplotlib.pyplot as plt import nir import numpy as np # for loading dataset import torch import spyx

/home/kade/Code/spyx/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

In [2]:

ng = nir.read("braille_noDelay_noBias_subtract_subgraph.nir")

ng = nir.read("braille_noDelay_noBias_subtract_subgraph.nir")

In [3]:

data = torch.load("ds_test.pt", weights_only=False)

data = torch.load("ds_test.pt", weights_only=False)

In [4]:

x, y = data.tensors

x, y = data.tensors

In [5]:

x = jnp.array(x)

x = jnp.array(x)

In [6]:

y = jnp.array(y)

y = jnp.array(y)

In [7]:

x.shape # --> 256 time

x.shape # --> 256 time

Out[7]:

(140, 256, 12)

In [8]:

ng.nodes.keys()

ng.nodes.keys()

Out[8]:

dict_keys(['fc1', 'fc2', 'input', 'lif1', 'lif2', 'output'])

In [9]:

ng.edges

ng.edges

Out[9]:

[('input', 'fc1'),
 ('fc1', 'lif1'),
 ('lif1', 'fc2'),
 ('fc2', 'lif2'),
 ('lif2', 'output')]

In [10]:

# Extract the input -> fc1 -> lif1 recurrent block as a standalone subgraph.
# The Output node must declare lif1's shape (read it from the graph so this
# works for every export), not the 7-class network output.
lif1_dim = int(np.asarray(ng.nodes["lif1"].output_type["output"]).reshape(-1)[0])
n_list = ["input", "fc1", "lif1"]
subgraph_nodes = {k: ng.nodes[k] for k in n_list}
subgraph_nodes["output"] = nir.Output(output_type={"output": np.array([lif1_dim])})
subgraph_edges = [("input", "fc1"), ("fc1", "lif1"), ("lif1", "output")]
subgraph = nir.NIRGraph(subgraph_nodes, subgraph_edges)

# Extract the input -> fc1 -> lif1 recurrent block as a standalone subgraph. # The Output node must declare lif1's shape (read it from the graph so this # works for every export), not the 7-class network output. lif1_dim = int(np.asarray(ng.nodes["lif1"].output_type["output"]).reshape(-1)[0]) n_list = ["input", "fc1", "lif1"] subgraph_nodes = {k: ng.nodes[k] for k in n_list} subgraph_nodes["output"] = nir.Output(output_type={"output": np.array([lif1_dim])}) subgraph_edges = [("input", "fc1"), ("fc1", "lif1"), ("lif1", "output")] subgraph = nir.NIRGraph(subgraph_nodes, subgraph_edges)

In [11]:

# (output shape already set when building the subgraph above)

# (output shape already set when building the subgraph above)

In [12]:

SNN, (a, b) = spyx.nir.from_nir(subgraph, x, dt=1e-4, return_all_states=True)

SNN, (a, b) = spyx.nir.from_nir(subgraph, x, dt=1e-4, return_all_states=True)

In [13]:

a.shape  # (T, B, ...) output spikes

a.shape # (T, B, ...) output spikes

Out[13]:

(140, 256, 40)

In [14]:

# `a` (output spikes) and `b` (states) were returned by from_nir above

# `a` (output spikes) and `b` (states) were returned by from_nir above

In [15]:

plt.figure(figsize=(12,4))
plt.imshow(x[0].T, aspect=4, interpolation="none")
plt.show()

plt.figure(figsize=(12,4)) plt.imshow(x[0].T, aspect=4, interpolation="none") plt.show()

In [16]:

a.shape

a.shape

Out[16]:

(140, 256, 40)

In [17]:

plt.figure(figsize=(12,4))
plt.imshow(a[0].T, aspect=4, interpolation="none")
plt.show()

plt.figure(figsize=(12,4)) plt.imshow(a[0].T, aspect=4, interpolation="none") plt.show()

In [18]:

np.save("spyx_activity_noDelay_noBias_subtract.npy", a[0])

np.save("spyx_activity_noDelay_noBias_subtract.npy", a[0])

zero partial network¶

In [19]:

ng = nir.read("braille_noDelay_bias_zero_subgraph.nir")

ng = nir.read("braille_noDelay_bias_zero_subgraph.nir")

In [20]:

# Extract the input -> fc1 -> lif1 recurrent block as a standalone subgraph.
# The Output node must declare lif1's shape (read it from the graph so this
# works for every export), not the 7-class network output.
lif1_dim = int(np.asarray(ng.nodes["lif1"].output_type["output"]).reshape(-1)[0])
n_list = ["input", "fc1", "lif1"]
subgraph_nodes = {k: ng.nodes[k] for k in n_list}
subgraph_nodes["output"] = nir.Output(output_type={"output": np.array([lif1_dim])})
subgraph_edges = [("input", "fc1"), ("fc1", "lif1"), ("lif1", "output")]
subgraph = nir.NIRGraph(subgraph_nodes, subgraph_edges)

# Extract the input -> fc1 -> lif1 recurrent block as a standalone subgraph. # The Output node must declare lif1's shape (read it from the graph so this # works for every export), not the 7-class network output. lif1_dim = int(np.asarray(ng.nodes["lif1"].output_type["output"]).reshape(-1)[0]) n_list = ["input", "fc1", "lif1"] subgraph_nodes = {k: ng.nodes[k] for k in n_list} subgraph_nodes["output"] = nir.Output(output_type={"output": np.array([lif1_dim])}) subgraph_edges = [("input", "fc1"), ("fc1", "lif1"), ("lif1", "output")] subgraph = nir.NIRGraph(subgraph_nodes, subgraph_edges)

In [21]:

# (output shape already set when building the subgraph above)

# (output shape already set when building the subgraph above)

In [22]:

SNN, (a, b) = spyx.nir.from_nir(subgraph, x, dt=1e-4, return_all_states=True)

SNN, (a, b) = spyx.nir.from_nir(subgraph, x, dt=1e-4, return_all_states=True)

In [23]:

a.shape  # (T, B, ...) output spikes

a.shape # (T, B, ...) output spikes

Out[23]:

(140, 256, 38)

In [24]:

# `a` (output spikes) and `b` (states) were returned by from_nir above

# `a` (output spikes) and `b` (states) were returned by from_nir above

In [25]:

plt.figure(figsize=(12,4))
plt.imshow(x[0].T, aspect=4, interpolation="none")
plt.show()

plt.figure(figsize=(12,4)) plt.imshow(x[0].T, aspect=4, interpolation="none") plt.show()

In [26]:

a.shape

a.shape

Out[26]:

(140, 256, 38)

In [27]:

plt.figure(figsize=(12,4))
plt.imshow(a[0].T, aspect=4, interpolation="none")
plt.show()

plt.figure(figsize=(12,4)) plt.imshow(a[0].T, aspect=4, interpolation="none") plt.show()

In [28]:

np.save("spyx_activity_noDelay_bias_zero.npy", a[0])

np.save("spyx_activity_noDelay_bias_zero.npy", a[0])

In [ ]:

Full networks¶

Run the complete imported RSNN (input -> fc1 -> lif1 -> fc2 -> lif2 -> output) over the Braille test set and score it with spyx.fn.integral_accuracy.

Expected result: this is an interoperability demo, not a tuned model. The imported noDelay / noBias / subtract-reset graph scores about 19% accuracy on the 7-way Braille test split (chance ~14%) — the point is that Spyx runs the NIR graph faithfully, not that this particular exported checkpoint classifies well.

In [29]:

ng = nir.read("braille_noDelay_noBias_subtract_subgraph.nir")

ng = nir.read("braille_noDelay_noBias_subtract_subgraph.nir")

In [30]:

SNN, (a, b) = spyx.nir.from_nir(ng, x, dt=1e-4, return_all_states=True)

SNN, (a, b) = spyx.nir.from_nir(ng, x, dt=1e-4, return_all_states=True)

In [31]:

# NNX: weights live in the reconstructed model, not a params dict.
# Sequential(Linear, RCuBaLIF): input weight on layer 0, recurrent on layer 1.
spyx_rec = SNN.layers[1].recurrent_w[...].T.flatten()
spyx_inp = SNN.layers[0].kernel[...].T.flatten()

# NNX: weights live in the reconstructed model, not a params dict. # Sequential(Linear, RCuBaLIF): input weight on layer 0, recurrent on layer 1. spyx_rec = SNN.layers[1].recurrent_w[...].T.flatten() spyx_inp = SNN.layers[0].kernel[...].T.flatten()

Optional cross-check against snnTorch¶

The original notebook re-imported the same NIR graph with snnTorch (from snntorch import import_nirtorch) and asserted that the input (fc1) and recurrent (lif1.recurrent) weight tensors matched Spyx's reconstruction to floating-point tolerance. That step needs the optional snntorch + nirtorch packages and PyTorch, so it is omitted from the rendered run. When both frameworks load the same NIR file the weights agree exactly (jnp.allclose(...) == True), which is the guarantee NIR is designed to provide: Spyx's importer reproduces the reference RSNN verbatim. The cell above shows how to reach those weights from the Spyx side (SNN.layers[0].kernel, SNN.layers[1].recurrent_w).

In [32]:

# `a` (output spikes) and `b` (states) were returned by from_nir above

# `a` (output spikes) and `b` (states) were returned by from_nir above

In [33]:

plt.figure(figsize=(12,4))
plt.imshow(a[0].T, aspect=10, interpolation="none")

plt.figure(figsize=(12,4)) plt.imshow(a[0].T, aspect=10, interpolation="none")

Out[33]:

<matplotlib.image.AxesImage at 0x731d34315ac0>

In [34]:

y

y

Out[34]:

Array([1, 3, 2, 2, 6, 1, 1, 3, 4, 5, 4, 0, 5, 5, 0, 2, 4, 3, 1, 2, 5, 2,
       4, 6, 2, 2, 4, 1, 4, 4, 1, 3, 2, 0, 4, 5, 1, 0, 3, 5, 1, 2, 0, 4,
       5, 4, 5, 6, 6, 1, 4, 5, 0, 2, 3, 4, 5, 0, 2, 5, 5, 5, 6, 5, 6, 4,
       1, 2, 6, 1, 0, 0, 6, 4, 0, 3, 3, 0, 1, 6, 2, 0, 3, 1, 0, 1, 2, 0,
       3, 0, 0, 0, 4, 6, 1, 3, 2, 5, 2, 6, 0, 5, 5, 0, 3, 1, 6, 6, 3, 2,
       4, 4, 6, 3, 6, 2, 2, 5, 3, 6, 2, 1, 3, 6, 5, 4, 5, 4, 1, 6, 3, 0,
       3, 6, 3, 1, 6, 4, 3, 1], dtype=int32)

In [35]:

Acc = spyx.fn.integral_accuracy()
acc, preds = Acc(a, y)
float(acc)

Acc = spyx.fn.integral_accuracy() acc, preds = Acc(a, y) float(acc)

Out[35]:

0.18571428954601288

In [36]:

np.save("spyx_accuracy_noDelay_noBias_subtract.npy", acc)

np.save("spyx_accuracy_noDelay_noBias_subtract.npy", acc)

Bias + Zero reset¶

The same pipeline for the noDelay / bias / zero-reset export. Different reset semantics (zero vs subtract) change the dynamics, so the score differs.

Expected result: the imported bias / zero-reset graph scores about 13% accuracy here — again the value is secondary to confirming the round-trip runs.

In [37]:

ng = nir.read("braille_noDelay_bias_zero_subgraph.nir")

ng = nir.read("braille_noDelay_bias_zero_subgraph.nir")

In [38]:

SNN, (a, b) = spyx.nir.from_nir(ng, x, dt=1e-4, return_all_states=True)

SNN, (a, b) = spyx.nir.from_nir(ng, x, dt=1e-4, return_all_states=True)

In [39]:

# NNX: weights live in the reconstructed model, not a params dict.
# Sequential(Linear, RCuBaLIF): input weight on layer 0, recurrent on layer 1.
spyx_rec = SNN.layers[1].recurrent_w[...].T.flatten()
spyx_inp = SNN.layers[0].kernel[...].T.flatten()

# NNX: weights live in the reconstructed model, not a params dict. # Sequential(Linear, RCuBaLIF): input weight on layer 0, recurrent on layer 1. spyx_rec = SNN.layers[1].recurrent_w[...].T.flatten() spyx_inp = SNN.layers[0].kernel[...].T.flatten()

Optional cross-check against snnTorch¶

The original notebook re-imported the same NIR graph with snnTorch (from snntorch import import_nirtorch) and asserted that the input (fc1) and recurrent (lif1.recurrent) weight tensors matched Spyx's reconstruction to floating-point tolerance. That step needs the optional snntorch + nirtorch packages and PyTorch, so it is omitted from the rendered run. When both frameworks load the same NIR file the weights agree exactly (jnp.allclose(...) == True), which is the guarantee NIR is designed to provide: Spyx's importer reproduces the reference RSNN verbatim. The cell above shows how to reach those weights from the Spyx side (SNN.layers[0].kernel, SNN.layers[1].recurrent_w).

In [40]:

# `a` (output spikes) and `b` (states) were returned by from_nir above

# `a` (output spikes) and `b` (states) were returned by from_nir above

In [41]:

plt.figure(figsize=(12,4))
plt.imshow(a[0].T, aspect=10, interpolation="none")

plt.figure(figsize=(12,4)) plt.imshow(a[0].T, aspect=10, interpolation="none")

Out[41]:

<matplotlib.image.AxesImage at 0x731d343aeae0>

In [42]:

y

y

Out[42]:

Array([1, 3, 2, 2, 6, 1, 1, 3, 4, 5, 4, 0, 5, 5, 0, 2, 4, 3, 1, 2, 5, 2,
       4, 6, 2, 2, 4, 1, 4, 4, 1, 3, 2, 0, 4, 5, 1, 0, 3, 5, 1, 2, 0, 4,
       5, 4, 5, 6, 6, 1, 4, 5, 0, 2, 3, 4, 5, 0, 2, 5, 5, 5, 6, 5, 6, 4,
       1, 2, 6, 1, 0, 0, 6, 4, 0, 3, 3, 0, 1, 6, 2, 0, 3, 1, 0, 1, 2, 0,
       3, 0, 0, 0, 4, 6, 1, 3, 2, 5, 2, 6, 0, 5, 5, 0, 3, 1, 6, 6, 3, 2,
       4, 4, 6, 3, 6, 2, 2, 5, 3, 6, 2, 1, 3, 6, 5, 4, 5, 4, 1, 6, 3, 0,
       3, 6, 3, 1, 6, 4, 3, 1], dtype=int32)

In [43]:

Acc = spyx.fn.integral_accuracy()
acc, preds = Acc(a, y)
float(acc)

Acc = spyx.fn.integral_accuracy() acc, preds = Acc(a, y) float(acc)

Out[43]:

0.12857143580913544

In [44]:

np.save("spyx_accuracy_noDelay_bias_zero.npy", acc)

np.save("spyx_accuracy_noDelay_bias_zero.npy", acc)