From ad3455b9982b3f8789861bd6a293cc7456214371 Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Fri, 14 Aug 2026 17:21:09 +0200 Subject: [PATCH 1/4] chore(ruff): Adapt to ruff@0.6.0 --- README.md | 2 +- docs/source/api_design.md | 5 +- docs/source/examples/index.md | 2 + docs/source/examples/lava/lava_to_nir.py | 262 +++++++++---------- docs/source/examples/lava/nir_to_lava.py | 266 +++++++++----------- docs/source/examples/nengo/nir-lorentz.py | 18 +- docs/source/examples/nengo/nir-nengo-lif.py | 15 +- docs/source/nirtorch/state.md | 13 +- docs/source/porting_nir.md | 12 +- docs/source/primitives.md | 14 +- docs/source/usage.md | 4 +- docs/source/working_with_nir.md | 9 +- docs/supported_primitives.py | 161 +++++++++--- nir/__init__.py | 11 +- nir/data_ir/__init__.py | 12 +- nir/data_ir/graph.py | 34 ++- nir/ir/__init__.py | 34 +-- nir/ir/conv.py | 37 +-- nir/ir/flatten.py | 14 +- nir/ir/graph.py | 123 +++------ nir/ir/linear.py | 12 +- nir/ir/neuron.py | 37 ++- nir/ir/node.py | 16 +- nir/ir/typing.py | 7 +- nir/ir/utils.py | 16 +- nir/serialization.py | 60 ++--- tests/test_data_ir.py | 9 +- tests/test_ir.py | 142 ++++------- tests/test_readwrite.py | 8 +- tests/test_utils.py | 8 +- 30 files changed, 628 insertions(+), 735 deletions(-) diff --git a/README.md b/README.md index b2c0cdc3..8f48af7d 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ Using NIR is typically a part of your favorite framework's workflow, but follows # Define a model my_model = ... # Save the model (source platform) -nir.write("my_graph.nir", my_model) +nir.write("my_graph.nir", my_model) # Load the model (target platform) imported_graph = nir.read("my_graph.nir") ``` diff --git a/docs/source/api_design.md b/docs/source/api_design.md index 1df7e0a9..de8f9745 100644 --- a/docs/source/api_design.md +++ b/docs/source/api_design.md @@ -37,10 +37,7 @@ That would look like this: `edges = [("my_node", "my_node")]`. In sum, a rudimentary, self-cyclic graph can be described in NIR as follows: ```python -NIRGraph( - nodes = {"my_node": MyNIRNode(np.array([...]))}, - edges = [("my_node", "my_node")] -) +NIRGraph(nodes={"my_node": MyNIRNode(np.array([...]))}, edges=[("my_node", "my_node")]) ``` ## Input and output types diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md index 9085927b..9d850ec9 100644 --- a/docs/source/examples/index.md +++ b/docs/source/examples/index.md @@ -18,6 +18,7 @@ Note that this requires you to provide a NIR model, so you need to find a way to The `nir.write` function takes two arguments: the file path and the model to write. ```python import nir + my_nir_graph = ... nir.write("my_graph.nir", my_model) ``` @@ -27,6 +28,7 @@ To read a model from a file, use the `nir.read` function. This function takes a single argument: the file path. ```python import nir + imported_graph = nir.read("my_graph.nir") ``` diff --git a/docs/source/examples/lava/lava_to_nir.py b/docs/source/examples/lava/lava_to_nir.py index 777a10e7..f6fab8c4 100644 --- a/docs/source/examples/lava/lava_to_nir.py +++ b/docs/source/examples/lava/lava_to_nir.py @@ -1,20 +1,19 @@ -import typing -import h5py -import re +import logging import pathlib -import nir +import re + +import h5py import numpy as np -import logging -from typing import Tuple -import collections.abc +import nir -PATH_TYPE = typing.Union[str, pathlib.Path] +PATH_TYPE = str | pathlib.Path +logger = logging.getLogger(__name__) -def read_cuba_lif(layer: h5py.Group, shape: Tuple[int] = None) -> nir.NIRNode: +def read_cuba_lif(layer: h5py.Group, shape: tuple[int] | None = None) -> nir.NIRNode: """Reads a CUBA LIF layer from a h5py.Group. - TODOs: + TODOs: - what if the layer is more than 1D? - handle scaleRho tauRho theta - support graded spikes @@ -22,49 +21,48 @@ def read_cuba_lif(layer: h5py.Group, shape: Tuple[int] = None) -> nir.NIRNode: - support other neuron types If the neuron model is not supported, a warning is logged and None is returned. """ - logging.debug(f"read_cuba_lif {layer['neuron']['type'][()]}") - - if 'gradedSpike' in layer['neuron']: - if layer['neuron']['gradedSpike'][()]: - logging.warning('graded spikes not supported') - - if layer['neuron']['type'][()] in [b'LOIHI', b'CUBA']: - if layer['neuron']['refDelay'][()] != 1: - logging.warning('refdelay not supported, setting to 1') - if layer['neuron']['vDecay'] == 0: - logging.warning('vDecay is 0, setting to inf') - if layer['neuron']['iDecay'] == 0: - logging.warning('iDecay is 0, setting to inf') - - # Lava-dl exports hardware fixed tipe to hdf5. For NIR, export in floating point. - dt = 1e-4 # This dt value is according to nir_to_lava.py script. https://github.com/neuromorphs/NIR/blob/main/paper/nir_to_lava.py#L67 - vdecay = layer['neuron']['vDecay'][()] / 4096 # Save the value in NIR as floating point - idecay = layer['neuron']['iDecay'][()] / 4096 # Save the value in NIR as floating point - thr = layer['neuron']['vThMant'][()] / 64 # Save the value in NIR as floating point - tau_mem = dt/float(vdecay) if vdecay != 0 else np.inf - tau_syn = dt/float(idecay) if idecay != 0 else np.inf - shape = layer['weight'].shape[0] - r = tau_mem/dt # no scaling of synaptic current - w_in = tau_syn/dt # no scaling of synaptic voltage + logger.debug(f"read_cuba_lif {layer['neuron']['type'][()]}") + + if "gradedSpike" in layer["neuron"] and layer["neuron"]["gradedSpike"][()]: + logger.warning("graded spikes not supported") + + if layer["neuron"]["type"][()] in [b"LOIHI", b"CUBA"]: + if layer["neuron"]["refDelay"][()] != 1: + logger.warning("refdelay not supported, setting to 1") + if layer["neuron"]["vDecay"] == 0: + logger.warning("vDecay is 0, setting to inf") + if layer["neuron"]["iDecay"] == 0: + logger.warning("iDecay is 0, setting to inf") + + # Lava-dl exports hardware fixed tipe to hdf5. For NIR, export in floating point. + dt = 1e-4 # This dt value is according to nir_to_lava.py script. https://github.com/neuromorphs/NIR/blob/main/paper/nir_to_lava.py#L67 + vdecay = layer["neuron"]["vDecay"][()] / 4096 # Save the value in NIR as floating point + idecay = layer["neuron"]["iDecay"][()] / 4096 # Save the value in NIR as floating point + thr = layer["neuron"]["vThMant"][()] / 64 # Save the value in NIR as floating point + tau_mem = dt / float(vdecay) if vdecay != 0 else np.inf + tau_syn = dt / float(idecay) if idecay != 0 else np.inf + shape = layer["weight"].shape[0] + r = tau_mem / dt # no scaling of synaptic current + w_in = tau_syn / dt # no scaling of synaptic voltage return nir.CubaLIF( tau_syn=np.full(shape, tau_syn), tau_mem=np.full(shape, tau_mem), - r=np.full(shape, r), - v_leak=np.full(shape, 0.), # currently no bias in Loihi's neurons + r=np.full(shape, r), + v_leak=np.full(shape, 0.0), # currently no bias in Loihi's neurons v_threshold=np.full(shape, thr), - w_in=np.full(shape,w_in), - v_reset=np.full(shape,0.) # LAVA-DL CUBA LiF always reset to 0 + w_in=np.full(shape, w_in), + v_reset=np.full(shape, 0.0), # LAVA-DL CUBA LiF always reset to 0 ) else: - logging.warning('currently only support for CUBA-LIF') - logging.error(f"no support for {layer['neuron']['type'][()]}") + logger.warning("currently only support for CUBA-LIF") + logger.error(f"no support for {layer['neuron']['type'][()]}") return None def read_node(network: h5py.Group) -> nir.NIRNode: """Read a graph from a HDF/conn5 file. - + TODOs: - support delay in convolutional layers """ @@ -73,132 +71,137 @@ def read_node(network: h5py.Group) -> nir.NIRNode: current_shape = None # need to sort keys as integers, otherwise does 1->10->2 - layer_keys = sorted(list(map(int, network.keys()))) + layer_keys = sorted(map(int, network.keys())) # iterate over layers for layer_idx_int in layer_keys: layer_idx = str(layer_idx_int) layer = network[layer_idx] - logging.info(f"--- Layer #{layer_idx}: {layer['type'][0].decode().upper()}") - logging.debug(f'current shape: {current_shape}') + logger.info(f"--- Layer #{layer_idx}: {layer['type'][0].decode().upper()}") + logger.debug(f"current shape: {current_shape}") - if layer['type'][0] == b'dense': + if layer["type"][0] == b"dense": # shape, type, neuron, inFeatures, outFeatures, weight, delay? - logging.debug(f'dense weights of shape {layer["weight"][:].shape}') + logger.debug(f"dense weights of shape {layer['weight'][:].shape}") # make sure weight matrix matches shape of previous layer if current_shape is None: - assert len(layer['weight'][:].shape) == 2, 'shape mismatch in dense' - current_shape = layer['weight'][:].shape[1] + assert len(layer["weight"][:].shape) == 2, "shape mismatch in dense" + current_shape = layer["weight"][:].shape[1] elif isinstance(current_shape, int): - assert current_shape == layer['weight'][:].shape[-1], 'shape mismatch in dense' + assert current_shape == layer["weight"][:].shape[-1], "shape mismatch in dense" else: - assert len(current_shape) == 1, 'shape mismatch in dense' - assert current_shape[0] == layer['weight'][:].shape[1], 'shape mismatch in dense' + assert len(current_shape) == 1, "shape mismatch in dense" + assert current_shape[0] == layer["weight"][:].shape[1], "shape mismatch in dense" # infer shape of current layer - assert len(layer['weight'][:].shape) in [1, 2], 'invalid dimension for dense layer' - current_shape = 1 if len(layer['weight'][:].shape) == 1 else layer['weight'][:].shape[0] + assert len(layer["weight"][:].shape) in [1, 2], "invalid dimension for dense layer" + current_shape = 1 if len(layer["weight"][:].shape) == 1 else layer["weight"][:].shape[0] # store the weight matrix (np.array, carrying over type) - if 'bias' in layer: - nodes.append(nir.Affine(weight=(layer['weight'][:] / 64), bias=layer['bias'][:])) + if "bias" in layer: + nodes.append(nir.Affine(weight=(layer["weight"][:] / 64), bias=layer["bias"][:])) else: - nodes.append(nir.Linear(weight=(layer['weight'][:] / 64))) # Save weights as floating point as well. + nodes.append( + nir.Linear(weight=(layer["weight"][:] / 64)) + ) # Save weights as floating point as well. # store the neuron group neuron = read_cuba_lif(layer) if neuron is None: - raise NotImplementedError('could not read neuron') + raise NotImplementedError("could not read neuron") nodes.append(neuron) # connect linear to neuron, neuron to next element - edges.append((len(nodes)-2, len(nodes)-1)) - edges.append((len(nodes)-1, len(nodes))) + edges.append((len(nodes) - 2, len(nodes) - 1)) + edges.append((len(nodes) - 1, len(nodes))) - elif layer['type'][0] == b'input': + elif layer["type"][0] == b"input": # iDecay, refDelay, scaleRho, tauRho, theta, type, vDecay, vThMant, wgtExp - current_shape = layer['shape'][:] - logging.warning('INPUT - not implemented yet') - logging.debug(f'keys: {layer.keys()}') - logging.debug(f'shape: {layer["shape"][:]}, bias: {layer["bias"][()]}, weight: {layer["weight"][()]}') - logging.debug(f'neuron keys: {", ".join(list(layer["neuron"].keys()))}') - - elif layer['type'][0] == b'flatten': + current_shape = layer["shape"][:] + logger.warning("INPUT - not implemented yet") + logger.debug(f"keys: {layer.keys()}") + logger.debug( + f"shape: {layer['shape'][:]}, bias: {layer['bias'][()]}, weight: {layer['weight'][()]}" + ) + logger.debug(f"neuron keys: {', '.join(list(layer['neuron'].keys()))}") + + elif layer["type"][0] == b"flatten": # shape, type - logging.debug(f"flattening shape (ignored): {layer['shape'][:]}") + logger.debug(f"flattening shape (ignored): {layer['shape'][:]}") # check last layer's size - assert len(nodes) > 0, 'flatten layer: must be preceded by a layer' - assert isinstance(current_shape, tuple), 'flatten layer: nothing to flatten' - last_node = nodes[-1] + assert len(nodes) > 0, "flatten layer: must be preceded by a layer" + assert isinstance(current_shape, tuple), "flatten layer: nothing to flatten" nodes.append(nir.Flatten(n_dims=1)) current_shape = int(np.prod(current_shape)) - edges.append((len(nodes)-1, len(nodes))) + edges.append((len(nodes) - 1, len(nodes))) - elif layer['type'][0] == b'conv': - # shape, type, neuron, inChannels, outChannels, kernelSize, stride, + elif layer["type"][0] == b"conv": + # shape, type, neuron, inChannels, outChannels, kernelSize, stride, # padding, dilation, groups, weight, delay? - weight = layer['weight'][:] - stride = layer['stride'][()] - pad = layer['padding'][()] - dil = layer['dilation'][()] - kernel_size = layer['kernelSize'][()] - in_channels = layer['inChannels'][()] - out_channels = layer['outChannels'][()] - logging.debug(f'stride {stride} padding {pad} dilation {dil} w {weight.shape}') + weight = layer["weight"][:] + stride = layer["stride"][()] + pad = layer["padding"][()] + dil = layer["dilation"][()] + kernel_size = layer["kernelSize"][()] + in_channels = layer["inChannels"][()] + out_channels = layer["outChannels"][()] + logger.debug(f"stride {stride} padding {pad} dilation {dil} w {weight.shape}") # infer shape of current layer - assert in_channels == current_shape[0], 'in_channels must match previous layer' + assert in_channels == current_shape[0], "in_channels must match previous layer" x_prev = current_shape[1] y_prev = current_shape[2] - x = (x_prev + 2*pad[0] - dil[0]*(kernel_size[0]-1) - 1) // stride[0] + 1 - y = (y_prev + 2*pad[1] - dil[1]*(kernel_size[1]-1) - 1) // stride[1] + 1 + x = (x_prev + 2 * pad[0] - dil[0] * (kernel_size[0] - 1) - 1) // stride[0] + 1 + y = (y_prev + 2 * pad[1] - dil[1] * (kernel_size[1] - 1) - 1) // stride[1] + 1 current_shape = (out_channels, x, y) # check for unsupported options - if layer['groups'][()] != 1: - logging.warning('groups not supported, setting to 1') - if 'delay' in layer: - logging.warning(f"delay=({layer['delay'][()]}) not supported, ignoring") + if layer["groups"][()] != 1: + logger.warning("groups not supported, setting to 1") + if "delay" in layer: + logger.warning(f"delay=({layer['delay'][()]}) not supported, ignoring") # store the conv matrix (np.array, carrying over type) - nodes.append(nir.Conv2d( - weight=layer['weight'][:], - bias=layer['bias'][:] if 'bias' in layer else None, - stride=stride, - padding=pad, - dilation=dil, - groups=layer['groups'][()] - )) + nodes.append( + nir.Conv2d( + weight=layer["weight"][:], + bias=layer["bias"][:] if "bias" in layer else None, + stride=stride, + padding=pad, + dilation=dil, + groups=layer["groups"][()], + ) + ) # store the neuron group neuron = read_cuba_lif(layer) if neuron is None: - raise NotImplementedError('could not read neuron') + raise NotImplementedError("could not read neuron") nodes.append(neuron) # connect conv to neuron group, neuron group to next element - edges.append((len(nodes)-2, len(nodes)-1)) - edges.append((len(nodes)-1, len(nodes))) + edges.append((len(nodes) - 2, len(nodes) - 1)) + edges.append((len(nodes) - 1, len(nodes))) - elif layer['type'][0] == b'average': + elif layer["type"][0] == b"average": # shape, type - logging.error('AVERAGE LAYER - not implemented yet') - raise NotImplementedError('average layer not implemented yet') + logger.error("AVERAGE LAYER - not implemented yet") + raise NotImplementedError("average layer not implemented yet") - elif layer['type'][0] == b'concat': + elif layer["type"][0] == b"concat": # shape, type, layers - logging.error('CONCAT LAYER - not implemented yet') - raise NotImplementedError('concat layer not implemented yet') + logger.error("CONCAT LAYER - not implemented yet") + raise NotImplementedError("concat layer not implemented yet") - elif layer['type'][0] == b'pool': + elif layer["type"][0] == b"pool": # shape, type, neuron, kernelSize, stride, padding, dilation, weight - logging.error('POOL LAYER - not implemented yet') - raise NotImplementedError('pool layer not implemented yet') + logger.error("POOL LAYER - not implemented yet") + raise NotImplementedError("pool layer not implemented yet") else: - logging.error('layer type not supported:', layer['type'][0]) + logger.error(f"layer type not supported: {layer['type'][0]}") # remove last edge (no next element) edges.pop(-1) @@ -210,50 +213,51 @@ def convert_to_nir(net_config: PATH_TYPE, path: PATH_TYPE) -> nir.NIRGraph: """Load a NIR from a HDF/conn5 file.""" with h5py.File(net_config, "r") as f: nir_graph = read_node(f["layer"]) - #nir_graph.input_type['input'] = nir_graph.input_type.pop('input_0') + # nir_graph.input_type['input'] = nir_graph.input_type.pop('input_0') nir_graph = normalize_nir_graph(nir_graph, to_bytes_in_edges=True) - #nir_graph.check_types() - nir_graph.edges[3]=(b'input',0) - nir_graph.edges[4]=(3,b'output') + # nir_graph.check_types() + nir_graph.edges[3] = (b"input", 0) + nir_graph.edges[4] = (3, b"output") nir.write(path, nir_graph) class Network: - def __init__(self, path: typing.Union[str, pathlib.Path]) -> None: + def __init__(self, path: str | pathlib.Path) -> None: nir_graph = nir.read(path) self.graph = nir_graph # TODO: implement the NIR -> Lava conversion - pass + def normalize_nir_graph(obj, to_bytes_in_edges=False, _in_edges=False): """ Recursively normalize all 'input_#'/'output_#' to 'input'/'output'. Only convert 'input'/'output' to bytes inside the 'edges' field if to_bytes_in_edges is True. """ + def norm(x, in_edges): # Normalize input/output with optional _# if isinstance(x, str): - if re.match(r'^input(_\d+)?$', x): - return b'input' if (to_bytes_in_edges and in_edges) else 'input' - if re.match(r'^output(_\d+)?$', x): - return b'output' if (to_bytes_in_edges and in_edges) else 'output' + if re.match(r"^input(_\d+)?$", x): + return b"input" if (to_bytes_in_edges and in_edges) else "input" + if re.match(r"^output(_\d+)?$", x): + return b"output" if (to_bytes_in_edges and in_edges) else "output" return x if isinstance(x, bytes): try: x_str = x.decode() - except Exception: + except (UnicodeDecodeError, AttributeError): return x - if re.match(r'^input(_\d+)?$', x_str): - return b'input' if (to_bytes_in_edges and in_edges) else 'input' - if re.match(r'^output(_\d+)?$', x_str): - return b'output' if (to_bytes_in_edges and in_edges) else 'output' + if re.match(r"^input(_\d+)?$", x_str): + return b"input" if (to_bytes_in_edges and in_edges) else "input" + if re.match(r"^output(_\d+)?$", x_str): + return b"output" if (to_bytes_in_edges and in_edges) else "output" return x return x if isinstance(obj, dict): new_dict = {} for k, v in obj.items(): - if k == 'edges' and isinstance(v, list): + if k == "edges" and isinstance(v, list): # Recursively process everything inside edges new_dict[k] = normalize_nir_graph(v, to_bytes_in_edges, _in_edges=True) else: @@ -265,10 +269,10 @@ def norm(x, in_edges): return tuple(normalize_nir_graph(v, to_bytes_in_edges, _in_edges) for v in obj) if isinstance(obj, set): return {normalize_nir_graph(v, to_bytes_in_edges, _in_edges) for v in obj} - if isinstance(obj, np.ndarray) and obj.dtype.kind == 'U': + if isinstance(obj, np.ndarray) and obj.dtype.kind == "U": return obj - if hasattr(obj, '__dict__'): + if hasattr(obj, "__dict__"): for k, v in vars(obj).items(): setattr(obj, k, normalize_nir_graph(v, to_bytes_in_edges, _in_edges)) return obj - return norm(obj, _in_edges) \ No newline at end of file + return norm(obj, _in_edges) diff --git a/docs/source/examples/lava/nir_to_lava.py b/docs/source/examples/lava/nir_to_lava.py index 411437b2..381b3b2a 100644 --- a/docs/source/examples/lava/nir_to_lava.py +++ b/docs/source/examples/lava/nir_to_lava.py @@ -4,19 +4,21 @@ Sharp edges: - in lava-dl, the current and voltage state is not automatically reset. must do this manually after every forward pass. """ -import nir -import nirtorch -import numpy as np + from dataclasses import dataclass -from functools import partial from enum import Enum -from lava.proc.lif.process import LIF -from lava.proc.dense.process import Dense -import lava.lib.dl.slayer as slayer +from functools import partial + +import nirtorch +import numpy as np import torch +from lava.lib.dl import slayer +from lava.proc.dense.process import Dense +from lava.proc.lif.process import LIF +import nir -LavaLibrary = Enum('LavaLibrary', 'Lava LavaDl') +LavaLibrary = Enum("LavaLibrary", "Lava LavaDl") @dataclass @@ -34,13 +36,14 @@ def __post_init__(self): # Lava helpers ############################## + def get_outport(lava_node): if isinstance(lava_node, Dense): return lava_node.a_out elif isinstance(lava_node, LIF): return lava_node.s_out else: - raise ValueError(f"Unknown node type: {type(lava_node)}") + raise TypeError(f"Unknown node type: {type(lava_node)}") def get_inport(lava_node): @@ -49,7 +52,7 @@ def get_inport(lava_node): elif isinstance(lava_node, LIF): return lava_node.a_in else: - raise ValueError(f"Unknown node type: {type(lava_node)}") + raise TypeError(f"Unknown node type: {type(lava_node)}") def _nir_node_to_lava(node: nir.NIRNode, import_config: ImportConfig): @@ -63,7 +66,6 @@ def _nir_node_to_lava(node: nir.NIRNode, import_config: ImportConfig): dv = dt / tau_mem vthr = node.v_threshold # * 10 # no current leak - tau_syn = None # 1/200 du = 1.0 # no current leak # correction for input weights correction = dt / node.tau @@ -77,12 +79,13 @@ def _nir_node_to_lava(node: nir.NIRNode, import_config: ImportConfig): w = (w * 256).astype(np.int32) lif = LIF( - shape=(1,), # u=0., # v=0., + shape=(1,), # u=0., # v=0., du=du, dv=dv, vth=vthr, - bias_mant=0, bias_exp=0, # no bias - name='lif' + bias_mant=0, + bias_exp=0, # no bias + name="lif", ) dense = Dense(weights=w) dense.a_out.connect(lif.a_in) @@ -90,20 +93,17 @@ def _nir_node_to_lava(node: nir.NIRNode, import_config: ImportConfig): elif isinstance(node, nir.Affine): w = node.weight - assert np.allclose(node.bias, 0.), "Non-zero bias not supported by Lava" + assert np.allclose(node.bias, 0.0), "Non-zero bias not supported by Lava" if import_config.fixed_pt: w = (w * 256).astype(np.int32) dense = Dense(weights=w) return dense - elif isinstance(node, nir.Input): - return None - - elif isinstance(node, nir.Output): + elif isinstance(node, (nir.Input, nir.Output)): return None else: - raise ValueError(f"Unknown node type: {type(node)}") + raise TypeError(f"Unknown node type: {type(node)}") def import_from_nir_to_lava(graph: nir.NIRGraph, import_config: ImportConfig): @@ -118,11 +118,10 @@ def import_from_nir_to_lava(graph: nir.NIRGraph, import_config: ImportConfig): start_nodes (List[str]): The start nodes in the graph end_nodes (List[str]): The end nodes in the graph """ - dt = import_config.dt - fixed_pt = import_config.fixed_pt lava_nodes = { - k: _nir_node_to_lava(n, import_config) for k, n in graph.nodes.items() + k: _nir_node_to_lava(n, import_config) + for k, n in graph.nodes.items() if not isinstance(n, (nir.Input, nir.Output)) } start_nodes = [] @@ -148,9 +147,10 @@ def import_from_nir_to_lava(graph: nir.NIRGraph, import_config: ImportConfig): # Lava-dl helpers ############################## + class Flatten(torch.nn.Module): def __init__(self, start_dim, end_dim): - super(Flatten, self).__init__() + super().__init__() self.start_dim = start_dim self.end_dim = end_dim @@ -159,9 +159,7 @@ def forward(self, x): end_dim = self.end_dim if self.end_dim == -1: end_dim -= 1 - elif self.end_dim == len(x.shape): - end_dim = -2 - elif self.end_dim == len(x.shape) - 1: + elif self.end_dim == len(x.shape) or self.end_dim == len(x.shape) - 1: end_dim = -2 # if end_dim != self.end_dim: # print(f'FLATTEN: changed end_dim from {self.start_dim} to {end_dim}') @@ -215,23 +213,22 @@ def _replace_rnn_subgraph_with_nirgraph(graph: nir.NIRGraph) -> tuple[nir.NIRGra n_subgraphs = 0 for edge1 in graph.edges: for edge2 in graph.edges: - if not edge1 == edge2: - if edge1[0] == edge2[1] and edge1[1] == edge2[0]: - lif_nk = edge1[0] - lif_n = graph.nodes[lif_nk] - w_nk = edge1[1] - w_n = graph.nodes[w_nk] - is_lif = isinstance(lif_n, (nir.LIF, nir.CubaLIF)) - is_dense = isinstance(w_n, (nir.Affine, nir.Linear)) - # check if the dense only connects to the LIF - w_out_nk = [e[1] for e in graph.edges if e[0] == w_nk] - w_in_nk = [e[0] for e in graph.edges if e[1] == w_nk] - is_rnn = len(w_out_nk) == 1 and len(w_in_nk) == 1 - # check if we found an RNN - if so, then parse it - if is_rnn and is_lif and is_dense: - print("creating rnn subgraph within nirgraph") - graph = _create_rnn_subgraph(graph, edge1[0], edge1[1]) - n_subgraphs += 1 + if edge1 != edge2 and edge1[0] == edge2[1] and edge1[1] == edge2[0]: + lif_nk = edge1[0] + lif_n = graph.nodes[lif_nk] + w_nk = edge1[1] + w_n = graph.nodes[w_nk] + is_lif = isinstance(lif_n, (nir.LIF, nir.CubaLIF)) + is_dense = isinstance(w_n, (nir.Affine, nir.Linear)) + # check if the dense only connects to the LIF + w_out_nk = [e[1] for e in graph.edges if e[0] == w_nk] + w_in_nk = [e[0] for e in graph.edges if e[1] == w_nk] + is_rnn = len(w_out_nk) == 1 and len(w_in_nk) == 1 + # check if we found an RNN - if so, then parse it + if is_rnn and is_lif and is_dense: + print("creating rnn subgraph within nirgraph") + graph = _create_rnn_subgraph(graph, edge1[0], edge1[1]) + n_subgraphs += 1 return graph, n_subgraphs @@ -250,18 +247,14 @@ def _parse_rnn_subgraph(graph: nir.NIRGraph) -> tuple[nir.NIRNode, nir.NIRNode, sub_nodes = graph.nodes.values() assert len(sub_nodes) == 4, "only 4-node RNN allowed in subgraph" try: - input_node = [n for n in sub_nodes if isinstance(n, nir.Input)][0] - output_node = [n for n in sub_nodes if isinstance(n, nir.Output)][0] - lif_node = [n for n in sub_nodes if isinstance(n, (nir.LIF, nir.CubaLIF))][0] - wrec_node = [n for n in sub_nodes if isinstance(n, (nir.Affine, nir.Linear))][0] + input_node = next(n for n in sub_nodes if isinstance(n, nir.Input)) + output_node = next(n for n in sub_nodes if isinstance(n, nir.Output)) + lif_node = next(n for n in sub_nodes if isinstance(n, (nir.LIF, nir.CubaLIF))) + wrec_node = next(n for n in sub_nodes if isinstance(n, (nir.Affine, nir.Linear))) except IndexError: raise ValueError("invalid RNN subgraph - could not find all required nodes") - lif_size = int( - list(input_node.input_type.values())[0][0] - ) # NOTE: needed for lava-dl - assert ( - lif_size == list(output_node.output_type.values())[0][0] - ), "output size mismatch" + lif_size = int(next(iter(input_node.input_type.values()))[0]) # NOTE: needed for lava-dl + assert lif_size == next(iter(output_node.output_type.values()))[0], "output size mismatch" assert lif_size == lif_node.v_threshold.size, "lif size mismatch (v_threshold)" assert lif_size == wrec_node.weight.shape[0], "w_rec shape mismatch" assert lif_size == wrec_node.weight.shape[1], "w_rec shape mismatch" @@ -293,60 +286,52 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): if debug_conv: print(f"Conv2d with weights of shape {node.weight.shape}:") print(f"\t{in_features} in, {out_features} out, kernel {kernel_size}") - print( - f"\tstride {node.stride}, padding {node.padding}, dilation {node.dilation}" - ) + print(f"\tstride {node.stride}, padding {node.padding}, dilation {node.dilation}") print(f"\tgroups {node.groups}") - conv_synapse_params = dict( - in_features=in_features, - out_features=out_features, - kernel_size=kernel_size, - stride=node.stride, - padding=node.padding, - dilation=node.dilation, - groups=node.groups, - weight_scale=1, - weight_norm=False, - pre_hook_fx=None, - ) + conv_synapse_params = { + "in_features": in_features, + "out_features": out_features, + "kernel_size": kernel_size, + "stride": node.stride, + "padding": node.padding, + "dilation": node.dilation, + "groups": node.groups, + "weight_scale": 1, + "weight_norm": False, + "pre_hook_fx": None, + } conv = slayer.synapse.Conv(**conv_synapse_params) conv.weight.data = torch.from_numpy(node.weight.reshape(conv.weight.shape)) return conv elif isinstance(node, nir.SumPool2d): if debug_pool: - print( - f"SumPool2d: kernel {node.kernel_size} pad {node.padding}, stride {node.stride}" - ) - pool_synapse_params = dict( - kernel_size=node.kernel_size, - stride=node.stride, - padding=node.padding, - dilation=1, - weight_scale=1, - weight_norm=False, - pre_hook_fx=None, - ) + print(f"SumPool2d: kernel {node.kernel_size} pad {node.padding}, stride {node.stride}") + pool_synapse_params = { + "kernel_size": node.kernel_size, + "stride": node.stride, + "padding": node.padding, + "dilation": 1, + "weight_scale": 1, + "weight_norm": False, + "pre_hook_fx": None, + } return slayer.synapse.Pool(**pool_synapse_params) elif isinstance(node, nir.IF): - assert ( - len(np.unique(node.v_threshold)) == 1 - ), "v_threshold must be the same for all neurons" - assert ( - len(np.unique(node.r)) == 1 - ), "resistance must be the same for all neurons" + assert len(np.unique(node.v_threshold)) == 1, "v_threshold must be the same for all neurons" + assert len(np.unique(node.r)) == 1, "resistance must be the same for all neurons" v_thr = np.unique(node.v_threshold)[0] resistance = np.unique(node.r)[0] v_thr_eff = v_thr * resistance * scale_v_thr if debug_if: print(f"IF with v_thr={v_thr}, R={resistance} -> eff. v_thr={v_thr_eff}") - cuba_neuron_params = dict( - threshold=v_thr_eff, - current_decay=1.0, - voltage_decay=0.0, - scale=4096, - ) + cuba_neuron_params = { + "threshold": v_thr_eff, + "current_decay": 1.0, + "voltage_decay": 0.0, + "scale": 4096, + } return slayer.neuron.cuba.Neuron(**cuba_neuron_params) # alif_neuron_params = dict( # threshold=v_thr_eff, threshold_step=0.0, scale=4096, @@ -375,15 +360,13 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): pre_hook_fx=None, ) dense.weight = torch.nn.Parameter( - data=torch.from_numpy(node.weight.reshape(dense.weight.shape)), - requires_grad=True + data=torch.from_numpy(node.weight.reshape(dense.weight.shape)), requires_grad=True ) dense.bias = torch.nn.Parameter( - data=torch.from_numpy(node.bias.reshape(node.weight.shape[0])), - requires_grad=True + data=torch.from_numpy(node.bias.reshape(node.weight.shape[0])), requires_grad=True ) return dense - + elif isinstance(node, nir.Linear): print("[WARNING] Linear layer not supported, using Dense instead") dense = slayer.synapse.Dense( @@ -394,17 +377,16 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): pre_hook_fx=None, ) dense.weight = torch.nn.Parameter( - data=torch.from_numpy(node.weight.reshape(dense.weight.shape)), + data=torch.from_numpy(node.weight.reshape(dense.weight.shape)), ) dense.bias = torch.nn.Parameter( - data=torch.zeros((node.weight.shape[0],)), - requires_grad=False + data=torch.zeros((node.weight.shape[0],)), requires_grad=False ) return dense elif isinstance(node, nir.CubaLIF): # TODO: figure out how to make the n_neurons dynamic - n_neurons = int(node.input_type['input'][0]) + n_neurons = int(node.input_type["input"][0]) # bias = node.v_leak * dt / node.tau_mem @@ -432,13 +414,13 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): weight_norm=False, pre_hook_fx=None, delay_shift=False, - neuron_params=dict( - threshold=np.unique(vthr)[0], - current_decay=np.unique(cur_decay)[0], - voltage_decay=np.unique(vol_decay)[0], - shared_param=True, - scale=scale, - ), + neuron_params={ + "threshold": np.unique(vthr)[0], + "current_decay": np.unique(cur_decay)[0], + "voltage_decay": np.unique(vol_decay)[0], + "shared_param": True, + "scale": scale, + }, ) # block.neuron.threshold_eps = 0.0 @@ -459,27 +441,17 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): elif isinstance(lif_node, nir.CubaLIF): # bias = lif_node.v_leak * dt / lif_node.tau_mem - assert np.allclose( - lif_node.v_leak, 0 - ), "v_leak not supported" # not yet in lava-dl? - assert np.allclose( - lif_node.r, lif_node.tau_mem / dt - ), "r not supported in CubaLIF" + assert np.allclose(lif_node.v_leak, 0), "v_leak not supported" # not yet in lava-dl? + assert np.allclose(lif_node.r, lif_node.tau_mem / dt), "r not supported in CubaLIF" cur_decay = dt / lif_node.tau_syn vol_decay = dt / lif_node.tau_mem w_scale = lif_node.w_in * (dt / lif_node.tau_syn) vthr = lif_node.v_threshold - assert ( - np.unique(cur_decay).size == 1 - ), "CubaLIF cur_decay must be same for all neurons" - assert ( - np.unique(vol_decay).size == 1 - ), "CubaLIF vol_decay must be same for all neurons" - assert ( - np.unique(vthr).size == 1 - ), "CubaLIF v_thr must be same for all neurons" + assert np.unique(cur_decay).size == 1, "CubaLIF cur_decay must be same for all neurons" + assert np.unique(vol_decay).size == 1, "CubaLIF vol_decay must be same for all neurons" + assert np.unique(vthr).size == 1, "CubaLIF v_thr must be same for all neurons" rnn_block = slayer.block.cuba.Recurrent( in_neurons=lif_size, @@ -488,13 +460,13 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): weight_norm=False, pre_hook_fx=None, delay_shift=False, - neuron_params=dict( - threshold=np.unique(vthr)[0], - current_decay=np.unique(cur_decay)[0], - voltage_decay=np.unique(vol_decay)[0], - shared_param=True, - scale=scale, - ), + neuron_params={ + "threshold": np.unique(vthr)[0], + "current_decay": np.unique(cur_decay)[0], + "voltage_decay": np.unique(vol_decay)[0], + "shared_param": True, + "scale": scale, + }, ) # rnn_block.neuron.threshold_eps = 0.0 @@ -504,21 +476,15 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): # TODO: make sure that dims match up print(f"[warning] scaling pre weights for w_in -> w_scale={w_scale[0]}") w_pre = w_pre * w_scale - rnn_block.input_synapse.weight = torch.nn.Parameter( - data=w_pre, requires_grad=True - ) + rnn_block.input_synapse.weight = torch.nn.Parameter(data=w_pre, requires_grad=True) wrec_shape = rnn_block.recurrent_synapse.weight.shape wrec = torch.from_numpy(wrec_node.weight).reshape(wrec_shape) - rnn_block.recurrent_synapse.weight = torch.nn.Parameter( - data=wrec, requires_grad=True - ) + rnn_block.recurrent_synapse.weight = torch.nn.Parameter(data=wrec, requires_grad=True) if isinstance(wrec_node, nir.Affine) and wrec_node.bias is not None: - bias = torch.from_numpy(wrec_node.bias).reshape((lif_size)) - rnn_block.recurrent_synapse.bias = torch.nn.Parameter( - data=bias, requires_grad=True - ) + bias = torch.from_numpy(wrec_node.bias).reshape(lif_size) + rnn_block.recurrent_synapse.bias = torch.nn.Parameter(data=bias, requires_grad=True) return rnn_block @@ -526,12 +492,12 @@ def _nir_node_to_lava_dl(node: nir.NIRNode, import_config: ImportConfig): raise NotImplementedError("LIF not implemented yet in lava-dl") else: - raise ValueError(f"Unknown node type: {type(node)}") + raise TypeError(f"Unknown node type: {type(node)}") class NIR2LavaDLNetwork(torch.nn.Module): def __init__(self, module_list, jens_order=False): - super(NIR2LavaDLNetwork, self).__init__() + super().__init__() self.blocks = torch.nn.ModuleList(module_list) def forward(self, spike): @@ -539,7 +505,7 @@ def forward(self, spike): if isinstance(block, torch.nn.Module): spike = block(spike) else: - raise Exception("Unknown block type") + raise TypeError("Unknown block type") return spike @@ -554,8 +520,6 @@ def get_next_node_key(node_key, edges): def import_from_nir_to_lava_dl(graph: nir.NIRGraph, import_config: ImportConfig, debug=False): # TODO (RNN addition): allow parsing of input and output nodes -> mapped to None - dt = import_config.dt - fixed_pt = import_config.fixed_pt # replace RNN subgraphs with NIRGraph nodes (i.e., subgraphs) graph, n_subgraphs = _replace_rnn_subgraph_with_nirgraph(graph) @@ -574,21 +538,21 @@ def import_from_nir_to_lava_dl(graph: nir.NIRGraph, import_config: ImportConfig, print(f"node {node_key}: {type(node).__name__}") if node_key == "output": continue - module_list.append( - _nir_node_to_lava_dl(node, import_config) - ) + module_list.append(_nir_node_to_lava_dl(node, import_config)) assert len(visited_node_keys) == len(graph.nodes), "not all nodes visited" return NIR2LavaDLNetwork(module_list) - + else: # found RNN subgraphs, need to use NIRTorch to parse the network net = nirtorch.load(graph, partial(_nir_node_to_lava_dl, import_config=import_config)) return net + ############################## # Main functions ############################## + def import_from_nir(graph: nir.NIRGraph, import_config: ImportConfig = None): if import_config is None: import_config = ImportConfig() @@ -598,4 +562,4 @@ def import_from_nir(graph: nir.NIRGraph, import_config: ImportConfig = None): elif import_config.library_preference == LavaLibrary.LavaDl: return import_from_nir_to_lava_dl(graph, import_config) else: - raise ValueError(f"Unknown library preference: {import_config.library_preference}") \ No newline at end of file + raise ValueError(f"Unknown library preference: {import_config.library_preference}") diff --git a/docs/source/examples/nengo/nir-lorentz.py b/docs/source/examples/nengo/nir-lorentz.py index 80c2aeca..324a0a1b 100644 --- a/docs/source/examples/nengo/nir-lorentz.py +++ b/docs/source/examples/nengo/nir-lorentz.py @@ -38,7 +38,7 @@ def nengo_to_nir(model): output = nengo.Node(lambda t, x: x, size_in=p.target.size_out) nengo.Connection(p.target, output, synapse=p.synapse) else: - raise Exception(f"Unhandled Probe {p}") + raise TypeError(f"Unhandled Probe {p}") sim = nengo.simulator.Simulator(model2) nengo2nir = {} @@ -119,9 +119,7 @@ def nengo_to_nir(model): if conn.synapse is not None: assert isinstance(conn.synapse, nengo.synapses.Lowpass) N = conn.size_out - ir = nir.LI( - tau=np.tile(conn.synapse.tau, N), r=np.tile(1, N), v_leak=np.tile(0, N) - ) + ir = nir.LI(tau=np.tile(conn.synapse.tau, N), r=np.tile(1, N), v_leak=np.tile(0, N)) nir_nodes.append(ir) nir_edges.append((source_index, len(nir_nodes) - 1)) source_index = len(nir_nodes) - 1 @@ -153,9 +151,7 @@ def nir_to_nengo(n): n_neurons=N, dimensions=1, label=f"LIF {i}", - neuron_type=nengo.RegularSpiking( - nengo.LIFRate(tau_rc=obj.tau[0], tau_ref=0) - ), + neuron_type=nengo.RegularSpiking(nengo.LIFRate(tau_rc=obj.tau[0], tau_ref=0)), # neuron_type=nengo.LIF(tau_rc=obj.tau[0], tau_ref=0), gain=np.ones(N), bias=np.zeros(N), @@ -182,7 +178,7 @@ def nir_to_nengo(n): None ) # because NIR spec doesn't tell me the size, I can't create this yet else: - raise Exception(f"Unknown NIR object: {obj}") + raise TypeError(f"Unknown NIR object: {obj}") for pre, post in n.edges: if nengo_map[post] is None: output = nengo.Node( @@ -194,11 +190,7 @@ def nir_to_nengo(n): synapse = filters.get(nengo_map[post], None) if nengo_map[pre].size_out != nengo_map[post].size_in: - print("Error") - print("pre", nengo_map[pre]) - print("post", nengo_map[post]) - 1 / 0 - + raise TypeError("Incompatible node sizes") else: nengo.Connection(nengo_map[pre], nengo_map[post], synapse=synapse) diff --git a/docs/source/examples/nengo/nir-nengo-lif.py b/docs/source/examples/nengo/nir-nengo-lif.py index b54c0dc0..9669f3f6 100644 --- a/docs/source/examples/nengo/nir-nengo-lif.py +++ b/docs/source/examples/nengo/nir-nengo-lif.py @@ -35,9 +35,7 @@ def nir_to_nengo(n, swap_linear_order=False): n_neurons=N, dimensions=1, label=f"LIF {i}", - neuron_type=nengo.RegularSpiking( - nengo.LIFRate(tau_rc=obj.tau[0], tau_ref=0) - ), + neuron_type=nengo.RegularSpiking(nengo.LIFRate(tau_rc=obj.tau[0], tau_ref=0)), # neuron_type=nengo.LIF(tau_rc=obj.tau[0], tau_ref=0), gain=np.ones(N), bias=np.zeros(N), @@ -56,7 +54,7 @@ def nir_to_nengo(n, swap_linear_order=False): if swap_linear_order: weights = weights.T w = nengo.Node( - lambda t, x, obj=obj: weights @ x + obj.bias, + lambda t, x, obj=obj, weights=weights: weights @ x + obj.bias, size_in=weights.shape[1], size_out=weights.shape[0], label=f"({weights.shape[0]}x{weights.shape[1]})", @@ -67,7 +65,7 @@ def nir_to_nengo(n, swap_linear_order=False): None ) # because NIR spec doesn't tell me the size, I can't create this yet else: - raise Exception(f"Unknown NIR object: {obj}") + raise TypeError(f"Unknown NIR object: {obj}") for pre, post in n.edges: if nengo_map[post] is None: output = nengo.Node( @@ -79,10 +77,7 @@ def nir_to_nengo(n, swap_linear_order=False): synapse = filters.get(nengo_map[post], None) if nengo_map[pre].size_out != nengo_map[post].size_in: - print("Error") - print("pre", nengo_map[pre]) - print("post", nengo_map[post]) - 1 / 0 + raise TypeError("Incompatible node sizes") else: nengo.Connection(nengo_map[pre], nengo_map[post], synapse=synapse) @@ -90,4 +85,4 @@ def nir_to_nengo(n, swap_linear_order=False): return model -model = nir_to_nengo(n, swap_linear_order=True) \ No newline at end of file +model = nir_to_nengo(n, swap_linear_order=True) diff --git a/docs/source/nirtorch/state.md b/docs/source/nirtorch/state.md index ab148885..1bb38891 100644 --- a/docs/source/nirtorch/state.md +++ b/docs/source/nirtorch/state.md @@ -108,16 +108,17 @@ NIRTorch uses **explicit** state management, which may be more cumbersome to wri class MyState: voltage: float + def stateful_function(data, state): # 1. Calculate a new voltage - new_voltage = ... + new_voltage = ... # 2. Calculate the function output - output = ... + output = ... # 3. Define a new state new_state = MyState(voltage=new_voltage) # 4. A tuple of (data, state) is returned # Note that the new state returned and the original remains unchanged - return output, new_state + return output, new_state ``` Once NIRTorch has parsed a NIR module into Torch modules (read more about that in the page about [To PyTorch: Interpreting NIR](#nirtorch_interpreting)), @@ -134,9 +135,9 @@ nir_weight = np.ones((2, 2)) nir_graph = nir.NIRGraph.from_list(nir.Linear(weight=nir_weight)) torch_module = nirtorch.nir_to_torch( - nir_graph=nir_graph, - node_map={} # We can leave this empty since we only - # use a linear layer which has a default mapping + nir_graph=nir_graph, + node_map={}, # We can leave this empty since we only + # use a linear layer which has a default mapping ) ## diff --git a/docs/source/porting_nir.md b/docs/source/porting_nir.md index 32fbbed9..f226adee 100644 --- a/docs/source/porting_nir.md +++ b/docs/source/porting_nir.md @@ -20,6 +20,7 @@ If your graph is stored in a file, you can load it using the `nir.read` function ```python import nir + my_graph = nir.read("path_to_my_graph.nir") ``` @@ -28,9 +29,9 @@ Note that the top-level graph may be recursive, so we recommend a recursive func Here's a simple example (with recursion): ```python - import nir + def parse_graph(graph: nir.NIRGraph): # Create a dictionary of nodes nodes = {} @@ -44,18 +45,17 @@ def parse_graph(graph: nir.NIRGraph): nodes[name] = MyPlatformAffine(node.weights, node.bias) elif isinstance(node, nir.Output): nodes[name] = MyPlatformOutput() - elif isinstance(node, nir.NIRGraph): # Recurse through subgraphs + elif isinstance(node, nir.NIRGraph): # Recurse through subgraphs nodes[name] = parse_graph(node) else: raise NotImplementedError(f"Node {node} not supported.") - + # Connect the nodes for edge in graph.edges: # Connect the nodes nodes[edge[0]].connect(nodes[edge[1]]) return nodes - ``` Matching the nodes to your primitives is the critical part here. @@ -78,6 +78,7 @@ Since several libraries are built on top of PyTorch, we provide default PyTorch ```python import nir, nirtorch + # Map nodes that are specific to your library # - nirtorch will map obvious nodes like `Input`, `Output`, `Affine`, `Conv2d` etc. # - but only if your parsing function do not return a module for that node @@ -85,7 +86,8 @@ def parse_module(node: nir.NIRNode) -> Optional[torch.nn.Module]: if isinstance(module, LIFBoxCell): return ... else: - return None # Return none to allow nirtorch to map the node + return None # Return none to allow nirtorch to map the node + # Interpret a NIR graph as a PyTorch module (`torch.nn.Module`) nir_graph = ... diff --git a/docs/source/primitives.md b/docs/source/primitives.md index 35a15975..01cefe1a 100644 --- a/docs/source/primitives.md +++ b/docs/source/primitives.md @@ -59,12 +59,9 @@ Note that a single node can be both an input and an output node. To clarify the dimensionality/input types of the input and output nodes, we require the user to specify the shape *and* name of the input, like so: ```python import numpy as np -nir.Input( - input_type = {"input": np.array([28, 28])} -) -nir.Output( - output_type = {"output": np.array([2])} -) + +nir.Input(input_type={"input": np.array([28, 28])}) +nir.Output(output_type={"output": np.array([2])}) ``` ## Metadata @@ -77,10 +74,7 @@ Here is an example of a metadata dictionary attached to a graph: ```python import nir -nir.NIRGraph( - ..., - metadata = {"some": "metadata", "info": 1} -) +nir.NIRGraph(..., metadata={"some": "metadata", "info": 1}) ``` diff --git a/docs/source/usage.md b/docs/source/usage.md index 61f04815..3e0a6a11 100644 --- a/docs/source/usage.md +++ b/docs/source/usage.md @@ -41,7 +41,7 @@ import sinabs from sinabs.backend.dynapcnn import DynapcnnNetwork # Convert NIR model to Sinabs -batch_size = ... # Define batch size to your liking +batch_size = ... # Define batch size to your liking sinabs.from_nir(nir_model, batch_size=batch_size) # Convert Sinabsmodel to chip-supported CNN dynapcnn_model = DynapcnnNetwork(sinabs_model, input_shape=sample_data.shape[-1]) @@ -67,6 +67,7 @@ To write this graph to file, you can use ```python import nir + nir.write(nir_model, "my_model.nir") ``` @@ -75,6 +76,7 @@ Reading a NIR file is similarly easy and will give you a graph object that you c ```python import nir + nir_model = nir.read("my_model.nir") ``` diff --git a/docs/source/working_with_nir.md b/docs/source/working_with_nir.md index adf54df7..5824aaf3 100644 --- a/docs/source/working_with_nir.md +++ b/docs/source/working_with_nir.md @@ -39,7 +39,8 @@ One example is the [`export_to_nir` function in snnTorch](https://snntorch.readt ```python import snntorch -my_snntorch_net = torch.nn.Sequential( ... ) + +my_snntorch_net = torch.nn.Sequential(...) nir_graph = snntorch.export_to_nir(my_snntorch_net) ``` @@ -50,6 +51,7 @@ Note that we provide a reference implementation in Python in the `nir` package, To read a model from a file, use the `nir.read` function with the path to the graph. ```python import nir + nir_graph = nir.read("my_graph.nir") ``` @@ -64,8 +66,8 @@ The topmost node is typically a [`NIRGraph` node](https://github.com/neuromorphs You can access those by the `.nodes` and `.edges` properties, respectively. ```python -nodes = nir_graph.nodes # A Dictionary of str -> nir.NIRNode -edges = nir_graph.edges # A List tuples (str, str) +nodes = nir_graph.nodes # A Dictionary of str -> nir.NIRNode +edges = nir_graph.edges # A List tuples (str, str) ``` Edges are pretty boring, because they just list the connections from one node to another. @@ -96,6 +98,7 @@ If you have a `NIRGraph` present, you can write it directly using the `nir.write The `nir.write` function takes two arguments: the file path and the model to write. ```python import nir + my_nir_graph = ... nir.write("my_graph.nir", my_model) ``` diff --git a/docs/supported_primitives.py b/docs/supported_primitives.py index 67241601..bdffe6f7 100644 --- a/docs/supported_primitives.py +++ b/docs/supported_primitives.py @@ -1,7 +1,8 @@ import re + import requests -from nir.ir import __all_ir as primitives +from nir.ir import __all_ir as primitives # Delete Input, Output, and NIRGraph from the list of primitives primitives = [p for p in primitives if p not in ["Input", "Output", "NIRGraph"]] @@ -9,44 +10,121 @@ # Fetch raw converter file from GitHub GITHUB_RAW_URLS = [ # (LibraryName, Direction, URL, function) - ("hxtorch", "from_nir", "https://raw.githubusercontent.com/electronicvisions/hxtorch/refs/heads/master/src/pyhxtorch/hxtorch/spiking/utils/from_nir.py", None), - ("hxtorch", "to_nir", "https://raw.githubusercontent.com/electronicvisions/hxtorch/refs/heads/master/src/pyhxtorch/hxtorch/spiking/utils/to_nir.py", None), - ("jaxsnn", "from_nir", "https://raw.githubusercontent.com/electronicvisions/jaxsnn/refs/heads/main/src/pyjaxsnn/jaxsnn/event/from_nir.py", None), - ("Lava", "from_nir", "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/paper/nir_to_lava.py", None), - ("Nengo", "from_nir", "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/docs/source/examples/nengo/nir-lorentz.py", "nir_to_nengo"), - ("Nengo", "to_nir", "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/docs/source/examples/nengo/nir-lorentz.py", "nengo_to_nir"), - ("Norse", "from_nir", "https://raw.githubusercontent.com/norse/norse/main/norse/torch/utils/import_nir.py", None), - ("Norse", "to_nir", "https://raw.githubusercontent.com/norse/norse/main/norse/torch/utils/export_nir.py", None), - ("rockpool", "from_nir", "https://raw.githubusercontent.com/synsense/rockpool/refs/heads/develop/rockpool/nn/modules/torch/nir.py", "_convert_nir_to_rockpool"), - ("rockpool", "to_nir", "https://raw.githubusercontent.com/synsense/rockpool/refs/heads/develop/rockpool/nn/modules/torch/nir.py", "_extract_rockpool_module"), - ("sinabs", "from_nir", "https://raw.githubusercontent.com/synsense/sinabs/refs/heads/develop/sinabs/nir.py", "_import_sinabs_module"), - ("sinabs", "to_nir", "https://raw.githubusercontent.com/synsense/sinabs/refs/heads/develop/sinabs/nir.py", "_extract_sinabs_module"), - ("snntorch", "from_nir", "https://raw.githubusercontent.com/jeshraghian/snntorch/refs/heads/master/snntorch/import_nir.py", None), - ("snntorch", "to_nir", "https://raw.githubusercontent.com/jeshraghian/snntorch/refs/heads/master/snntorch/export_nir.py", None), - ("SpiNNaker2", "from_nir", "https://gitlab.com/spinnaker2/py-spinnaker2/-/raw/main/src/spinnaker2/s2_nir.py?ref_type=heads", None), - ("Spyx", "from_nir", "https://raw.githubusercontent.com/kmheckel/spyx/refs/heads/main/src/spyx/nir.py", "_nir_node_to_spyx_node"), - ("Spyx", "to_nir", "https://raw.githubusercontent.com/kmheckel/spyx/refs/heads/main/src/spyx/nir.py", "to_nir"), + ( + "hxtorch", + "from_nir", + "https://raw.githubusercontent.com/electronicvisions/hxtorch/refs/heads/master/src/pyhxtorch/hxtorch/spiking/utils/from_nir.py", + None, + ), + ( + "hxtorch", + "to_nir", + "https://raw.githubusercontent.com/electronicvisions/hxtorch/refs/heads/master/src/pyhxtorch/hxtorch/spiking/utils/to_nir.py", + None, + ), + ( + "jaxsnn", + "from_nir", + "https://raw.githubusercontent.com/electronicvisions/jaxsnn/refs/heads/main/src/pyjaxsnn/jaxsnn/event/from_nir.py", + None, + ), + ( + "Lava", + "from_nir", + "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/paper/nir_to_lava.py", + None, + ), + ( + "Nengo", + "from_nir", + "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/docs/source/examples/nengo/nir-lorentz.py", + "nir_to_nengo", + ), + ( + "Nengo", + "to_nir", + "https://raw.githubusercontent.com/neuromorphs/NIR/refs/heads/main/docs/source/examples/nengo/nir-lorentz.py", + "nengo_to_nir", + ), + ( + "Norse", + "from_nir", + "https://raw.githubusercontent.com/norse/norse/main/norse/torch/utils/import_nir.py", + None, + ), + ( + "Norse", + "to_nir", + "https://raw.githubusercontent.com/norse/norse/main/norse/torch/utils/export_nir.py", + None, + ), + ( + "rockpool", + "from_nir", + "https://raw.githubusercontent.com/synsense/rockpool/refs/heads/develop/rockpool/nn/modules/torch/nir.py", + "_convert_nir_to_rockpool", + ), + ( + "rockpool", + "to_nir", + "https://raw.githubusercontent.com/synsense/rockpool/refs/heads/develop/rockpool/nn/modules/torch/nir.py", + "_extract_rockpool_module", + ), + ( + "sinabs", + "from_nir", + "https://raw.githubusercontent.com/synsense/sinabs/refs/heads/develop/sinabs/nir.py", + "_import_sinabs_module", + ), + ( + "sinabs", + "to_nir", + "https://raw.githubusercontent.com/synsense/sinabs/refs/heads/develop/sinabs/nir.py", + "_extract_sinabs_module", + ), + ( + "snntorch", + "from_nir", + "https://raw.githubusercontent.com/jeshraghian/snntorch/refs/heads/master/snntorch/import_nir.py", + None, + ), + ( + "snntorch", + "to_nir", + "https://raw.githubusercontent.com/jeshraghian/snntorch/refs/heads/master/snntorch/export_nir.py", + None, + ), + ( + "SpiNNaker2", + "from_nir", + "https://gitlab.com/spinnaker2/py-spinnaker2/-/raw/main/src/spinnaker2/s2_nir.py?ref_type=heads", + None, + ), + ( + "Spyx", + "from_nir", + "https://raw.githubusercontent.com/kmheckel/spyx/refs/heads/main/src/spyx/nir.py", + "_nir_node_to_spyx_node", + ), + ( + "Spyx", + "to_nir", + "https://raw.githubusercontent.com/kmheckel/spyx/refs/heads/main/src/spyx/nir.py", + "to_nir", + ), ] converter_contents = {} for lib_name, direction, url, function in GITHUB_RAW_URLS: if function is None: key = f"{lib_name}_{direction}" - try: - converter_contents[key] = requests.get(url).text - except Exception as e: - converter_contents[key] = "" - print(f"Failed to fetch {url}: {e}") + converter_contents[key] = requests.get(url).text else: key = f"{lib_name}_{direction}" - try: - response = requests.get(url) - pattern = rf"def {function}\s*\(.*?\):(.*?)(?=\ndef |\Z)" - match = re.search(pattern, response.text, re.DOTALL) - converter_contents[key] = match.group(0) - except Exception as e: - converter_contents[key] = "" - print(f"Failed to fetch {url}: {e}") + response = requests.get(url) + pattern = rf"def {function}\s*\(.*?\):(.*?)(?=\ndef |\Z)" + match = re.search(pattern, response.text, re.DOTALL) + converter_contents[key] = match.group(0) # Check which primitives are supported in each library/direction supported = {} @@ -55,8 +133,16 @@ for lib_name, _, _, _ in GITHUB_RAW_URLS: from_key = f"{lib_name}_from_nir" to_key = f"{lib_name}_to_nir" - from_supported = "ir.{})".format(name) in converter_contents.get(from_key, "") or "ir.{}(".format(name) in converter_contents.get(from_key, "") or "ir.{}:".format(name) in converter_contents.get(from_key, "") - to_supported = "ir.{})".format(name) in converter_contents.get(to_key, "") or "ir.{}(".format(name) in converter_contents.get(to_key, "") or "ir.{}:".format(name) in converter_contents.get(to_key, "") + from_supported = ( + f"ir.{name})" in converter_contents.get(from_key, "") + or f"ir.{name}(" in converter_contents.get(from_key, "") + or f"ir.{name}:" in converter_contents.get(from_key, "") + ) + to_supported = ( + f"ir.{name})" in converter_contents.get(to_key, "") + or f"ir.{name}(" in converter_contents.get(to_key, "") + or f"ir.{name}:" in converter_contents.get(to_key, "") + ) if from_supported and to_supported: supported[name][lib_name] = "⟷" elif from_supported: @@ -85,7 +171,6 @@ - `←`: Supported for conversion to NIR - `⟷`: Supported for both conversion directions (to and from NIR) Please note that this list is generated automatically and may not be entirely accurate. -
""" full_md = static_md + "\n\n" + dynamic_md @@ -96,13 +181,15 @@ # Generate enumeration of supported primitives for each library for lib in libs: - with open(f"docs/tmp/examples/{lib.lower()}/supported_primitives.md", "w", encoding="utf-8") as f: + with open( + f"docs/tmp/examples/{lib.lower()}/supported_primitives.md", "w", encoding="utf-8" + ) as f: support_to_nir = any(supported[p][lib] in ["←", "⟷"] for p in primitives) support_from_nir = any(supported[p][lib] in ["→", "⟷"] for p in primitives) lib_md = f"### Supported Primitives in {lib}\n\n" if support_to_nir: - lib_md += "This library supports conversion of the following nodes to NIR:" + lib_md += "This library supports conversion of the following nodes to NIR:" else: lib_md += "This library does not support conversion of any nodes to NIR." for primitive in primitives: @@ -110,7 +197,7 @@ lib_md += f"\n- {primitive}" if support_from_nir: - lib_md += "\n\nThis library supports conversion of the following nodes from NIR:" + lib_md += "\n\nThis library supports conversion of the following nodes from NIR:" else: lib_md += "This library does not support conversion of any nodes from NIR." for primitive in primitives: diff --git a/nir/__init__.py b/nir/__init__.py index ec6f6d92..c43e2e4e 100644 --- a/nir/__init__.py +++ b/nir/__init__.py @@ -3,7 +3,8 @@ Documentation: https://nnir.readthedocs.io """ -from importlib.metadata import version as metadata_version, PackageNotFoundError +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as metadata_version try: __version__ = version = metadata_version("nir") @@ -12,10 +13,10 @@ # package is not installed pass -from . import ir, data_ir -from .ir import * # noqa: F403 +from . import data_ir, ir +from .data_ir import * +from .ir import * from .ir import typing # noqa: F401 -from .data_ir import * # noqa: F403 -from .serialization import read, write, read_data, write_data +from .serialization import read, read_data, write, write_data __all__ = ir.__all__ + data_ir.__all__ + ["read", "write", "read_data", "write_data"] diff --git a/nir/data_ir/__init__.py b/nir/data_ir/__init__.py index 1c037388..12a5e1ef 100644 --- a/nir/data_ir/__init__.py +++ b/nir/data_ir/__init__.py @@ -1,15 +1,15 @@ from .graph import ( EventData, - ValuedEventData, - TimeGriddedData, - NIRNodeData, NIRGraphData, + NIRNodeData, + TimeGriddedData, + ValuedEventData, ) __all__ = [ - "TimeGriddedData", "EventData", - "ValuedEventData", - "NIRNodeData", "NIRGraphData", + "NIRNodeData", + "TimeGriddedData", + "ValuedEventData", ] diff --git a/nir/data_ir/graph.py b/nir/data_ir/graph.py index a5275549..653735ff 100644 --- a/nir/data_ir/graph.py +++ b/nir/data_ir/graph.py @@ -1,7 +1,9 @@ from __future__ import annotations + from dataclasses import dataclass -from typing import Dict, Union + import numpy as np + from nir.ir import NIRGraph, NIRNode @@ -26,8 +28,7 @@ class TimeGriddedData: def __post_init__(self): if not isinstance(self.data, np.ndarray) or self.data.ndim != 3: raise ValueError( - "Data must be of shape (n_samples, n_time_steps, n_neurons)" - "and of type np.ndarray" + "Data must be of shape (n_samples, n_time_steps, n_neurons)and of type np.ndarray" ) def __getitem__(self, idx): @@ -125,7 +126,8 @@ def n_samples(self): return self.idx.shape[0] def to_time_gridded( - self, dt: float # pylint: disable=invalid-name + self, + dt: float, # pylint: disable=invalid-name ) -> TimeGriddedData: """ Arguments @@ -134,14 +136,12 @@ def to_time_gridded( Time step size. """ n_time_steps = int(self.t_max / dt) - discrete_data = np.zeros( - (self.n_samples, n_time_steps, self.n_neurons), dtype=bool - ) + discrete_data = np.zeros((self.n_samples, n_time_steps, self.n_neurons), dtype=bool) for sample in range(self.n_samples): valid_spikes = self.idx[sample] != -1 valid_times = self.time[sample][valid_spikes] - steps = np.floor((valid_times / dt)).astype(int) + steps = np.floor(valid_times / dt).astype(int) neurons = self.idx[sample][valid_spikes] discrete_data[sample, steps, neurons] = True return TimeGriddedData(discrete_data, dt) @@ -193,7 +193,7 @@ def to_time_gridded( for sample in range(n_samples): valid_spikes = self.idx[sample] != -1 valid_times = self.time[sample][valid_spikes] - steps = np.floor((valid_times / dt)).astype(int) + steps = np.floor(valid_times / dt).astype(int) neurons = self.idx[sample][valid_spikes] value = self.value[sample][valid_spikes] discrete_data[sample, steps, neurons] = value @@ -213,13 +213,11 @@ class NIRNodeData: Dictionary of observables for a NIRNode. """ - observables: Dict[str, Union[EventData, TimeGriddedData]] + observables: dict[str, EventData | TimeGriddedData] def __post_init__(self): if not isinstance(self.observables, dict): - raise TypeError( - "observables must be a dictionary of EventData or TimeGriddedData" - ) + raise TypeError("observables must be a dictionary of EventData or TimeGriddedData") def __getitem__(self, idx): return self.observables[idx] @@ -232,9 +230,7 @@ def check_observables(self, node: NIRNode): Check that the shapes of the observables match the node's output shapes """ output_shape = node.output_type["output"] - if not all(obs.n_neurons == output_shape for obs in self.observables.values()): - return False - return True + return all(obs.n_neurons == output_shape for obs in self.observables.values()) @dataclass @@ -249,7 +245,7 @@ class NIRGraphData: Dictionary of NIRNodeData or NIRGraphData for a NIRGraph. """ - nodes: Dict[str, Union["NIRGraphData", NIRNodeData]] + nodes: dict[str, NIRGraphData | NIRNodeData] def __post_init__(self): if not isinstance(self.nodes, dict): @@ -278,6 +274,4 @@ def check_nodes(self, graph: NIRGraph): if not isinstance(graph_node, NIRNode): raise TypeError(f"Node {key} is not a NIRNode in the NIRGraph") if not node.check_observables(graph_node): - raise ValueError( - f"Observables for node {key} do not match the NIRNode" - ) + raise ValueError(f"Observables for node {key} do not match the NIRNode") diff --git a/nir/ir/__init__.py b/nir/ir/__init__.py index ee9bb863..6437a796 100644 --- a/nir/ir/__init__.py +++ b/nir/ir/__init__.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import Any from .conv import Conv1d, Conv2d from .delay import Delay @@ -48,7 +48,7 @@ def str2NIRNode(type: str) -> NIRNode: return globals()[type] -def dict2NIRNode(data_dict: Dict[str, Any]) -> NIRNode: +def dict2NIRNode(data_dict: dict[str, Any]) -> NIRNode: """Assume data_dict["type"] exist and correspond to a subclass of NIRNode. Other items should match fields in the corresponding NIRNode subclass, unless @@ -61,35 +61,35 @@ def dict2NIRNode(data_dict: Dict[str, Any]) -> NIRNode: # we could do this, but ruff complains # __all__ = __all_ir + ["str2NIRNode", "dict2NIRNode"] __all__ = [ + "IF", + "LI", + "LIF", + # linear + "Affine", + # pooling + "AvgPool2d", # conv "Conv1d", "Conv2d", + # neuron + "CubaLI", + "CubaLIF", # delay "Delay", # flatten "Flatten", + "I", # graph "Input", - "NIRGraph", - "Output", - # linear - "Affine", "Linear", - "Scale", - # neuron - "CubaLI", - "CubaLIF", - "I", - "IF", - "LI", - "LIF", + "NIRGraph", # node "NIRNode", - # pooling - "AvgPool2d", + "Output", + "Scale", "SumPool2d", # surrogate_gradient "Threshold", - "str2NIRNode", "dict2NIRNode", + "str2NIRNode", ] diff --git a/nir/ir/conv.py b/nir/ir/conv.py index 8df05dde..8b93d425 100644 --- a/nir/ir/conv.py +++ b/nir/ir/conv.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np @@ -34,28 +33,24 @@ class Conv1d(NIRNode): :type bias: np.ndarray """ - input_shape: Optional[int] # N + input_shape: int | None # N weight: np.ndarray # Weight C_out * C_in * N stride: int # Stride - padding: Union[int, str] # Padding + padding: int | str # Padding dilation: int # Dilation groups: int # Groups bias: np.ndarray # Bias C_out def __post_init__(self): if isinstance(self.padding, str) and self.padding not in ["same", "valid"]: - raise ValueError( - f"padding must be 'same', 'valid', or int, not {self.padding}" - ) + raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}") if self.input_shape is None: # leave input and output types undefined self.input_type = {"input": None} self.output_type = {"output": None} else: # infer input and output types from input_shape - self.input_type = { - "input": np.array([self.weight.shape[1], self.input_shape]) - } + self.input_type = {"input": np.array([self.weight.shape[1], self.input_shape])} output_shape = calculate_conv_output( self.input_shape, self.padding, @@ -63,9 +58,7 @@ def __post_init__(self): self.weight.shape[2], self.stride, ) - self.output_type = { - "output": np.array([self.weight.shape[0], *output_shape]) - } + self.output_type = {"output": np.array([self.weight.shape[0], *output_shape])} @dataclass(eq=False) @@ -95,19 +88,17 @@ class Conv2d(NIRNode): :type bias: np.ndarray """ - input_shape: Optional[Tuple[int, int]] # N_x, N_y + input_shape: tuple[int, int] | None # N_x, N_y weight: np.ndarray # Weight C_out * C_in * W_x * W_y - stride: Union[int, Tuple[int, int]] # Stride - padding: Union[int, Tuple[int, int], str] # Padding - dilation: Union[int, Tuple[int, int]] # Dilation + stride: int | tuple[int, int] # Stride + padding: int | tuple[int, int] | str # Padding + dilation: int | tuple[int, int] # Dilation groups: int # Groups bias: np.ndarray # Bias C_out def __post_init__(self): if isinstance(self.padding, str) and self.padding not in ["same", "valid"]: - raise ValueError( - f"padding must be 'same', 'valid', or int, not {self.padding}" - ) + raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}") if isinstance(self.padding, int): self.padding = (self.padding, self.padding) if isinstance(self.stride, int): @@ -120,9 +111,7 @@ def __post_init__(self): self.output_type = {"output": None} else: # infer input and output types from input_shape - self.input_type = { - "input": np.array([self.weight.shape[1], *self.input_shape]) - } + self.input_type = {"input": np.array([self.weight.shape[1], *self.input_shape])} output_shape = calculate_conv_output( self.input_shape, self.padding, @@ -130,6 +119,4 @@ def __post_init__(self): self.weight.shape[2], self.stride, ) - self.output_type = { - "output": np.array([self.weight.shape[0], *output_shape]) - } + self.output_type = {"output": np.array([self.weight.shape[0], *output_shape])} diff --git a/nir/ir/flatten.py b/nir/ir/flatten.py index 6c088ca4..23c29402 100644 --- a/nir/ir/flatten.py +++ b/nir/ir/flatten.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Any, Dict +from typing import Any import numpy as np @@ -35,18 +35,14 @@ def __post_init__(self): } # make sure input and output shape are valid if np.prod(self.input_type["input"]) != np.prod(self.output_type["output"]): - raise ValueError( - "input and output shape must have same number of elements" - ) + raise ValueError("input and output shape must have same number of elements") - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: ret = super().to_dict() ret["input_type"] = self.input_type["input"] return ret @classmethod - def from_dict(cls, node: Dict[str, Any]): - node["input_type"] = { - "input": node["input_type"] if "input_type" in node else None - } + def from_dict(cls, node: dict[str, Any]): + node["input_type"] = {"input": node.get("input_type", None)} return super().from_dict(node) diff --git a/nir/ir/graph.py b/nir/ir/graph.py index 41aea1e6..625d37a9 100644 --- a/nir/ir/graph.py +++ b/nir/ir/graph.py @@ -1,6 +1,6 @@ from collections import Counter from dataclasses import dataclass -from typing import Any, Dict +from typing import Any import numpy as np @@ -39,7 +39,7 @@ def __init__( self, nodes: Nodes, edges: Edges, - metadata: Dict[str, Any] = dict, + metadata: dict[str, Any] = dict, type_check: bool = True, ): self.nodes = nodes @@ -57,30 +57,24 @@ def __init__( @property def inputs(self): - return { - name: node for name, node in self.nodes.items() if isinstance(node, Input) - } + return {name: node for name, node in self.nodes.items() if isinstance(node, Input)} @property def outputs(self): - return { - name: node for name, node in self.nodes.items() if isinstance(node, Output) - } + return {name: node for name, node in self.nodes.items() if isinstance(node, Output)} @staticmethod def from_list(*nodes: NIRNode, type_check: bool = True) -> "NIRGraph": """Create a sequential graph from a list of nodes by labelling them after indices.""" - if len(nodes) > 0 and ( - isinstance(nodes[0], list) or isinstance(nodes[0], tuple) - ): + if len(nodes) > 0 and isinstance(nodes[0], (list, tuple)): nodes = [*nodes[0]] def unique_node_name(node, counts): basename = node.__class__.__name__.lower() id = counts[basename] - name = f"{basename}{f'_{id}' if id>0 else ''}" + name = f"{basename}{f'_{id}' if id > 0 else ''}" counts[basename] += 1 return name @@ -116,32 +110,24 @@ def __post_init__(self): self.metadata = {} def _update_input_output_types(self): - input_node_keys = [ - k for k, node in self.nodes.items() if isinstance(node, Input) - ] + input_node_keys = [k for k, node in self.nodes.items() if isinstance(node, Input)] self.input_type = ( - { - node_key: self.nodes[node_key].input_type["input"] - for node_key in input_node_keys - } + {node_key: self.nodes[node_key].input_type["input"] for node_key in input_node_keys} if len(input_node_keys) > 0 else None ) - output_node_keys = [ - k for k, node in self.nodes.items() if isinstance(node, Output) - ] + output_node_keys = [k for k, node in self.nodes.items() if isinstance(node, Output)] self.output_type = { - node_key: self.nodes[node_key].output_type["output"] - for node_key in output_node_keys + node_key: self.nodes[node_key].output_type["output"] for node_key in output_node_keys } - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: ret = super().to_dict() ret["nodes"] = {k: n.to_dict() for k, n in self.nodes.items()} return ret @classmethod - def from_dict(cls, kwargs: Dict[str, Any]) -> "NIRGraph": + def from_dict(cls, kwargs: dict[str, Any]) -> "NIRGraph": from . import dict2NIRNode kwargs_local = kwargs.copy() # Copy the input to avoid overwriting attributes @@ -151,20 +137,16 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "NIRGraph": assert "edges" in kwargs, "The incoming dictionary must hade a 'edges' entry" # Assert that the type is well-formed if "type" in kwargs: - assert ( - kwargs["type"] == "NIRGraph" - ), "You are calling NIRGraph.from_dict with a different type " + assert kwargs["type"] == "NIRGraph", ( + "You are calling NIRGraph.from_dict with a different type " + ) f"{type}. Either remove the entry or use .from_dict, such as Input.from_dict" kwargs_local["type"] = "NIRGraph" - kwargs_local["nodes"] = { - k: dict2NIRNode(n) for k, n in kwargs_local["nodes"].items() - } + kwargs_local["nodes"] = {k: dict2NIRNode(n) for k, n in kwargs_local["nodes"].items()} # h5py deserializes edges into a numpy array of type bytes and dtype=object, # hence using ensure_str here - kwargs_local["edges"] = [ - (ensure_str(a), ensure_str(b)) for a, b in kwargs_local["edges"] - ] + kwargs_local["edges"] = [(ensure_str(a), ensure_str(b)) for a, b in kwargs_local["edges"]] return super().from_dict(kwargs_local) def validate_structure(self): @@ -240,16 +222,14 @@ def check_types(self): # make sure the type values match up if len(pre_node.output_type.keys()) == 1: - post_input_type = list(post_node.input_type.values())[0] - pre_output_type = list(pre_node.output_type.values())[0] + post_input_type = next(iter(post_node.input_type.values())) + pre_output_type = next(iter(pre_node.output_type.values())) if not np.array_equal(post_input_type, pre_output_type): pre_repr = f"{edge[0]}.output: {pre_output_type}" post_repr = f"{edge[1]}.input: {post_input_type}" raise ValueError(f"type mismatch: {pre_repr} -> {post_repr}") else: - raise NotImplementedError( - "multiple input/output types not supported yet" - ) + raise NotImplementedError("multiple input/output types not supported yet") return True def infer_types(self): @@ -274,7 +254,7 @@ def infer_types(self): destination_nodes = {edge[1] for edge in self.edges} root_nodes = all_node_keys - destination_nodes - new_nodes: Dict[str, NIRNode] = {} + new_nodes: dict[str, NIRNode] = {} new_edges: Edges = [] for node_key in root_nodes: @@ -316,7 +296,7 @@ def infer_types(self): self.edges.extend(new_edges) # Start type inference from input nodes - ready = [e for e in self.edges if e[0] in self.inputs.keys()] + ready = [e for e in self.edges if e[0] in self.inputs] if len(ready) == 0: raise ValueError( "Failed to start type inference: No input nodes found. " @@ -325,7 +305,7 @@ def infer_types(self): "or disable type checking (`type_check=False`)." ) - seen = set([e[0] for e in ready]) + seen = {e[0] for e in ready} while len(ready) > 0: pre_key, post_key = ready.pop() pre_node = self.nodes[pre_key] @@ -350,26 +330,18 @@ def infer_types(self): if undef_post_input_type: # define post input_type to be the same as pre output_type post_node.input_type = { - k.replace("output", "input"): v - for k, v in pre_node.output_type.items() + k.replace("output", "input"): v for k, v in pre_node.output_type.items() } elif type_mismatch: # set post input_type to be the same as pre output_type - pre_repr = ( - f"{pre_key}.output: {np.array(list(pre_node.output_type.values()))}" - ) - post_repr = ( - f"{post_key}.input: {np.array(list(post_node.input_type.values()))}" - ) - raise ValueError( - f"Type inference error: type mismatch: {pre_repr} -> {post_repr}" - ) + pre_repr = f"{pre_key}.output: {np.array(list(pre_node.output_type.values()))}" + post_repr = f"{post_key}.input: {np.array(list(post_node.input_type.values()))}" + raise ValueError(f"Type inference error: type mismatch: {pre_repr} -> {post_repr}") # make sure that output nodes have output_type = input_type if isinstance(post_node, Output): post_node.output_type = { - k.replace("input", "output"): v - for k, v in post_node.input_type.items() + k.replace("input", "output"): v for k, v in post_node.input_type.items() } # check if post output_type needs to be defined @@ -378,7 +350,7 @@ def infer_types(self): ) if undef_post_output_type: # define post output_type - if isinstance(post_node, Conv1d) or isinstance(post_node, Conv2d): + if isinstance(post_node, (Conv1d, Conv2d)): if isinstance(post_node, Conv1d): post_node.input_shape = post_node.input_type["input"][1] else: @@ -393,20 +365,7 @@ def infer_types(self): output_type = np.array([post_node.weight.shape[0], *output_shape]) post_node.output_type = {"output": output_type} - elif isinstance(post_node, SumPool2d): - output_shape = calculate_conv_output( - pre_node.output_type["output"][1:], - post_node.padding, - 1, - post_node.kernel_size, - post_node.stride, - ) - output_type = np.array( - [post_node.input_type["input"][0], *output_shape] - ) - post_node.output_type = {"output": output_type} - - elif isinstance(post_node, AvgPool2d): + elif isinstance(post_node, (SumPool2d, AvgPool2d)): output_shape = calculate_conv_output( pre_node.output_type["output"][1:], post_node.padding, @@ -414,9 +373,7 @@ def infer_types(self): post_node.kernel_size, post_node.stride, ) - output_type = np.array( - [post_node.input_type["input"][0], *output_shape] - ) + output_type = np.array([post_node.input_type["input"][0], *output_shape]) post_node.output_type = {"output": output_type} elif isinstance(post_node, Flatten): @@ -429,9 +386,7 @@ def infer_types(self): } n_inputs = np.prod(post_node.input_type["input"]) n_outputs = np.prod(post_node.output_type["output"]) - assert ( - n_inputs == n_outputs - ), "Flatten must not change the number of elements" + assert n_inputs == n_outputs, "Flatten must not change the number of elements" seen.add(post_key) ready += [e for e in self.edges if e[0] == post_key and e[1] not in seen] @@ -449,7 +404,7 @@ def infer_types(self): "or disable type checking (`type_check=False`)." ) - new_nodes: Dict[str, NIRNode] = {} + new_nodes: dict[str, NIRNode] = {} new_edges: Edges = [] for node_key in leaf_nodes: @@ -503,13 +458,13 @@ def __post_init__(self): self.input_type = parse_shape_argument(self.input_type, "input") self.output_type = {"output": self.input_type["input"]} - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: ret = super().to_dict() ret["shape"] = self.input_type["input"] return ret @classmethod - def from_dict(cls, node: Dict[str, Any]) -> "NIRNode": + def from_dict(cls, node: dict[str, Any]) -> "NIRNode": node["input_type"] = {"input": node["shape"]} del node["shape"] return super().from_dict(node) @@ -530,13 +485,13 @@ def __post_init__(self): self.output_type = parse_shape_argument(self.output_type, "output") self.input_type = {"input": self.output_type["output"]} - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: ret = super().to_dict() ret["shape"] = self.output_type["output"] return ret @classmethod - def from_dict(cls, node: Dict[str, Any]) -> "NIRNode": + def from_dict(cls, node: dict[str, Any]) -> "NIRNode": node["output_type"] = {"output": node["shape"]} del node["shape"] return super().from_dict(node) @@ -556,11 +511,11 @@ class Identity(NIRNode): def __post_init__(self): self.output_type = self.input_type - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: ret = super().to_dict() ret["shape"] = self.output_type["output"] return ret @classmethod - def from_dict(cls, node: Dict[str, Any]) -> "NIRNode": + def from_dict(cls, node: dict[str, Any]) -> "NIRNode": return super().from_dict(node) diff --git a/nir/ir/linear.py b/nir/ir/linear.py index 2b4ea4b0..ffcf50af 100644 --- a/nir/ir/linear.py +++ b/nir/ir/linear.py @@ -24,13 +24,9 @@ class Affine(NIRNode): def __post_init__(self): assert len(self.weight.shape) >= 2, "Weight must be at least 2D" self.input_type = { - "input": np.array( - self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T) - ) - } - self.output_type = { - "output": np.array(self.weight.shape[:-2] + (self.weight.shape[-2],)) + "input": np.array(self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T)) } + self.output_type = {"output": np.array(self.weight.shape[:-2] + (self.weight.shape[-2],))} @dataclass(eq=False) @@ -46,9 +42,7 @@ class Linear(NIRNode): def __post_init__(self): assert len(self.weight.shape) >= 2, "Weight must be at least 2D" self.input_type = { - "input": np.array( - self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T) - ) + "input": np.array(self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T)) } self.output_type = {"output": self.weight.shape[:-2] + (self.weight.shape[-2],)} diff --git a/nir/ir/neuron.py b/nir/ir/neuron.py index 7bda56c0..573f332d 100644 --- a/nir/ir/neuron.py +++ b/nir/ir/neuron.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import Any import numpy as np @@ -34,12 +34,9 @@ class CubaLI(NIRNode): w_in: np.ndarray = 1.0 # Input current weight def __post_init__(self): - assert ( - self.tau_syn.shape - == self.tau_mem.shape - == self.r.shape - == self.v_leak.shape - ), "All parameters must have the same shape" + assert self.tau_syn.shape == self.tau_mem.shape == self.r.shape == self.v_leak.shape, ( + "All parameters must have the same shape" + ) # If w_in is a scalar, make it an array of same shape as v_leak self.w_in = np.ones_like(self.v_leak) * self.w_in self.input_type = {"input": np.array(self.v_leak.shape)} @@ -86,7 +83,7 @@ class CubaLIF(NIRNode): r: np.ndarray # Resistance v_leak: np.ndarray # Leak voltage v_threshold: np.ndarray # Firing threshold - v_reset: Optional[np.ndarray] = None # Reset potential + v_reset: np.ndarray | None = None # Reset potential w_in: np.ndarray = 1.0 # Input current weight def __post_init__(self): @@ -106,14 +103,14 @@ def __post_init__(self): self.output_type = {"output": np.array(self.v_threshold.shape)} @classmethod - def from_dict(cls, kwargs: Dict[str, Any]) -> "CubaLIF": + def from_dict(cls, kwargs: dict[str, Any]) -> "CubaLIF": if "v_reset" not in kwargs: kwargs["v_reset"] = np.zeros_like(kwargs["v_threshold"]) return super().from_dict(kwargs) @dataclass(eq=False) -class I(NIRNode): # noqa: E742 +class I(NIRNode): r"""Integrator. The integrator neuron model is defined by the following equation: @@ -153,19 +150,19 @@ class IF(NIRNode): r: np.ndarray # Resistance v_threshold: np.ndarray # Firing threshold - v_reset: Optional[np.ndarray] = None # Reset potential + v_reset: np.ndarray | None = None # Reset potential def __post_init__(self): if self.v_reset is None: self.v_reset = np.zeros_like(self.v_threshold) - assert ( - self.r.shape == self.v_threshold.shape == self.v_reset.shape - ), "All parameters must have the same shape" + assert self.r.shape == self.v_threshold.shape == self.v_reset.shape, ( + "All parameters must have the same shape" + ) self.input_type = {"input": np.array(self.r.shape)} self.output_type = {"output": np.array(self.r.shape)} @classmethod - def from_dict(cls, kwargs: Dict[str, Any]) -> "IF": + def from_dict(cls, kwargs: dict[str, Any]) -> "IF": if "v_reset" not in kwargs: kwargs["v_reset"] = np.zeros_like(kwargs["v_threshold"]) return super().from_dict(kwargs) @@ -190,9 +187,9 @@ class LI(NIRNode): v_leak: np.ndarray # Leak voltage def __post_init__(self): - assert ( - self.tau.shape == self.r.shape == self.v_leak.shape - ), "All parameters must have the same shape" + assert self.tau.shape == self.r.shape == self.v_leak.shape, ( + "All parameters must have the same shape" + ) self.input_type = {"input": np.array(self.r.shape)} self.output_type = {"output": np.array(self.r.shape)} @@ -231,7 +228,7 @@ class LIF(NIRNode): r: np.ndarray # Resistance v_leak: np.ndarray # Leak voltage v_threshold: np.ndarray # Firing threshold - v_reset: Optional[np.ndarray] = None # Reset potential + v_reset: np.ndarray | None = None # Reset potential def __post_init__(self): if self.v_reset is None: @@ -247,7 +244,7 @@ def __post_init__(self): self.output_type = {"output": np.array(self.r.shape)} @classmethod - def from_dict(cls, kwargs: Dict[str, Any]) -> "LIF": + def from_dict(cls, kwargs: dict[str, Any]) -> "LIF": if "v_reset" not in kwargs: kwargs["v_reset"] = np.zeros_like(kwargs["v_threshold"]) return super().from_dict(kwargs) diff --git a/nir/ir/node.py b/nir/ir/node.py index 7ab97010..d19817a4 100644 --- a/nir/ir/node.py +++ b/nir/ir/node.py @@ -1,6 +1,6 @@ from abc import ABC from dataclasses import asdict, dataclass, field -from typing import Any, Dict +from typing import Any import numpy as np @@ -18,9 +18,9 @@ class NIRNode(ABC): # keyword argument. All three are keyword-only so that subclasses can add # positional fields without running into the "non-default argument follows # default argument" ordering error. (Requires Python 3.10+.) - input_type: Dict[str, np.ndarray] = field(init=False, kw_only=True) - output_type: Dict[str, np.ndarray] = field(init=False, kw_only=True) - metadata: Dict[str, Any] = field(default_factory=dict, kw_only=True) + input_type: dict[str, np.ndarray] = field(init=False, kw_only=True) + output_type: dict[str, np.ndarray] = field(init=False, kw_only=True) + metadata: dict[str, Any] = field(default_factory=dict, kw_only=True) def __init__(self) -> None: raise AttributeError("NIRNode does not have a default constructor.") @@ -28,12 +28,12 @@ def __init__(self) -> None: def __eq__(self, other): return self is other - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """Serialize into a dictionary.""" ret = asdict(self) - if "input_type" in ret.keys(): + if "input_type" in ret: del ret["input_type"] - if "output_type" in ret.keys(): + if "output_type" in ret: del ret["output_type"] # Note: The customization below won't be automatically done recursively for nested NIRNode. # Therefore, classes with nested NIRNode e.g. NIRGraph must implement its own to_dict @@ -42,7 +42,7 @@ def to_dict(self) -> Dict[str, Any]: return ret @classmethod - def from_dict(cls, kwargs: Dict[str, Any]) -> "NIRNode": + def from_dict(cls, kwargs: dict[str, Any]) -> "NIRNode": assert kwargs["type"] == cls.__name__ kwargs = kwargs.copy() # Local scope del kwargs["type"] diff --git a/nir/ir/typing.py b/nir/ir/typing.py index d662c2ad..9e870838 100644 --- a/nir/ir/typing.py +++ b/nir/ir/typing.py @@ -1,12 +1,11 @@ -from typing import Dict, List, Tuple import numpy as np from .node import NIRNode # Nodes are uniquely named computational units -Nodes = Dict[str, "NIRNode"] +Nodes = dict[str, "NIRNode"] # Edges map one node id to another via the identity -Edges = List[Tuple[str, str]] +Edges = list[tuple[str, str]] # Types is a dict mapping strings to tensor shapes -Types = Dict[str, np.ndarray] +Types = dict[str, np.ndarray] diff --git a/nir/ir/utils.py b/nir/ir/utils.py index 3e864fdd..69bfcca6 100644 --- a/nir/ir/utils.py +++ b/nir/ir/utils.py @@ -1,4 +1,4 @@ -from typing import Sequence, Union +from collections.abc import Sequence import numpy as np @@ -18,11 +18,11 @@ def parse_shape_argument(x: Types, key: str): def calculate_conv_output( - input_shape: Union[int, Sequence[int]], - padding: Union[int, str, Sequence[int]], - dilation: Union[int, Sequence[int]], - kernel_size: Union[int, Sequence[int]], - stride: Union[int, Sequence[int]], + input_shape: int | Sequence[int], + padding: int | str | Sequence[int], + dilation: int | Sequence[int], + kernel_size: int | Sequence[int], + stride: int | Sequence[int], ) -> Sequence[int]: """Calculates the output for a single dimension of a convolutional layer. https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html#torch.nn.Conv1d @@ -87,7 +87,7 @@ def calc_flatten_output(input_shape: Sequence[int], start_dim: int, end_dim: int ) -def _index_tuple(tuple: Union[int, Sequence[int]], index: int) -> np.ndarray: +def _index_tuple(tuple: int | Sequence[int], index: int) -> np.ndarray: """If the input is a tuple/array, index it. Otherwise, return it as-is. @@ -105,7 +105,7 @@ def _index_tuple(tuple: Union[int, Sequence[int]], index: int) -> np.ndarray: raise TypeError(f"tuple must be int or np.ndarray, not {type(tuple)}") -def ensure_str(a: Union[str, bytes]) -> str: +def ensure_str(a: str | bytes) -> str: if isinstance(a, bytes): return a.decode("utf8") elif isinstance(a, str): diff --git a/nir/serialization.py b/nir/serialization.py index 1f734c07..28cc5473 100644 --- a/nir/serialization.py +++ b/nir/serialization.py @@ -1,6 +1,6 @@ import io import pathlib -from typing import Any, Dict, Union +from typing import Any import h5py import numpy as np @@ -8,28 +8,24 @@ import nir -def _read_metadata(node: Any) -> Dict[str, Any]: - if "metadata" in node.keys(): +def _read_metadata(node: Any) -> dict[str, Any]: + if "metadata" in node: return {"metadata": {k: v[()] for k, v in node["metadata"].items()}} else: return {} -def try_byte_to_str(a: Union[bytes, Any]) -> Union[str, Any]: +def try_byte_to_str(a: bytes | Any) -> str | Any: return a.decode("utf8") if isinstance(a, bytes) else a def read_node(node: Any) -> nir.NIRNode: """Read a graph from a HDF5 file.""" if node["type"][()] == b"Affine": - return nir.Affine( - weight=node["weight"][()], bias=node["bias"][()], **_read_metadata(node) - ) + return nir.Affine(weight=node["weight"][()], bias=node["bias"][()], **_read_metadata(node)) elif node["type"][()] == b"Conv1d": return nir.Conv1d( - input_shape=( - node["input_shape"][()] if "input_shape" in node.keys() else None - ), + input_shape=(node["input_shape"][()] if "input_shape" in node else None), weight=node["weight"][()], stride=node["stride"][()], padding=node["padding"][()], @@ -40,9 +36,7 @@ def read_node(node: Any) -> nir.NIRNode: ) elif node["type"][()] == b"Conv2d": return nir.Conv2d( - input_shape=( - node["input_shape"][()] if "input_shape" in node.keys() else None - ), + input_shape=(node["input_shape"][()] if "input_shape" in node else None), weight=node["weight"][()], stride=node["stride"][()], padding=node["padding"][()], @@ -70,9 +64,7 @@ def read_node(node: Any) -> nir.NIRNode: return nir.Flatten( start_dim=node["start_dim"][()], end_dim=node["end_dim"][()], - input_type={ - "input": node["input_type"][()] if "input_type" in node.keys() else None - }, + input_type={"input": node["input_type"][()] if "input_type" in node else None}, **_read_metadata(node), ) elif node["type"][()] == b"I": @@ -81,17 +73,13 @@ def read_node(node: Any) -> nir.NIRNode: return nir.IF( r=node["r"][()], v_reset=( - node["v_reset"][()] - if "v_reset" in node - else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], **_read_metadata(node), ) elif node["type"][()] == b"Input": - return nir.Input( - input_type={"input": node["shape"][()]}, **_read_metadata(node) - ) + return nir.Input(input_type={"input": node["shape"][()]}, **_read_metadata(node)) elif node["type"][()] == b"LI": return nir.LI( tau=node["tau"][()], @@ -107,9 +95,7 @@ def read_node(node: Any) -> nir.NIRNode: r=node["r"][()], v_leak=node["v_leak"][()], v_reset=( - node["v_reset"][()] - if "v_reset" in node - else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], **_read_metadata(node), @@ -130,9 +116,7 @@ def read_node(node: Any) -> nir.NIRNode: r=node["r"][()], v_leak=node["v_leak"][()], v_reset=( - node["v_reset"][()] - if "v_reset" in node - else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], w_in=node["w_in"][()], @@ -145,9 +129,7 @@ def read_node(node: Any) -> nir.NIRNode: **_read_metadata(node), ) elif node["type"][()] == b"Output": - return nir.Output( - output_type={"output": node["shape"][()]}, **_read_metadata(node) - ) + return nir.Output(output_type={"output": node["shape"][()]}, **_read_metadata(node)) elif node["type"][()] == b"Scale": return nir.Scale(scale=node["scale"][()], **_read_metadata(node)) elif node["type"][()] == b"Threshold": @@ -156,7 +138,7 @@ def read_node(node: Any) -> nir.NIRNode: raise ValueError(f"Unknown unit type: {node['type'][()]}") -def hdf2dict(node: Any) -> Dict[str, Any]: +def hdf2dict(node: Any) -> dict[str, Any]: ret = {} def read_hdf_to_dict(node, data_dict): @@ -173,7 +155,7 @@ def read_hdf_to_dict(node, data_dict): return ret -def read(filename: Union[str, pathlib.Path], type_check: bool = True) -> nir.NIRGraph: +def read(filename: str | pathlib.Path, type_check: bool = True) -> nir.NIRGraph: """Load a NIR from a HDF/conn5 file. Attempts to read a NIRGraph from a file and pass in the key-value parameters to the corresponding NIR nodes. @@ -198,7 +180,7 @@ def read(filename: Union[str, pathlib.Path], type_check: bool = True) -> nir.NIR return nir.dict2NIRNode(data_dict) -def read_version(filename: Union[str, pathlib.Path]) -> str: +def read_version(filename: str | pathlib.Path) -> str: """Reads the filename of a given NIR file, and raises an exception if the version does not exist in the file. @@ -210,7 +192,7 @@ def read_version(filename: Union[str, pathlib.Path]) -> str: def write( - filename: Union[str, pathlib.Path, io.RawIOBase], + filename: str | pathlib.Path | io.RawIOBase, graph: nir.NIRNode, compression: str = "gzip", compression_opts: Any = None, @@ -237,7 +219,7 @@ def write( def write_recursive(group: h5py.Group, node: dict) -> None: for k, v in node.items(): if k == "metadata": - if not v == {}: # Skip metadata if empty + if v != {}: # Skip metadata if empty write_recursive(group.create_group(k), v) elif isinstance(v, str): group.create_dataset(k, data=v, dtype=h5py.string_dtype()) @@ -334,7 +316,7 @@ def _read_graph_data(group: h5py.Group) -> nir.NIRGraphData: def write_data( - filename: Union[str, pathlib.Path, io.RawIOBase], + filename: str | pathlib.Path | io.RawIOBase, graph_data: nir.NIRGraphData, compression: str = "gzip", compression_opts: Any = None, @@ -410,11 +392,11 @@ def _write_node_data(group: h5py.Group, node_data: nir.NIRNodeData): else: raise TypeError(f"Unsupported observable type: {type(obs)}") - def _write_graph_data(group: h5py.Group, node: dict) -> None: + def _write_graph_data(group: h5py.Group, graph: nir.NIRGraphData) -> None: group.attrs["__type__"] = "NIRGraphData" nodes_group = group.create_group("nodes") - for name, node in graph_data.nodes.items(): + for name, node in graph.nodes.items(): g = nodes_group.create_group(name) if isinstance(node, nir.NIRNodeData): _write_node_data(g, node) diff --git a/tests/test_data_ir.py b/tests/test_data_ir.py index 8b012e64..ca59b428 100644 --- a/tests/test_data_ir.py +++ b/tests/test_data_ir.py @@ -1,4 +1,5 @@ import numpy as np + import nir @@ -7,7 +8,7 @@ def test_generate_time_gridded_data(): dt = 0.1 gridded = nir.TimeGriddedData(spikes, dt) node = nir.NIRNodeData({"spikes": gridded}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["spikes"].data, spikes) assert graph.nodes["node"].observables["spikes"].dt == dt @@ -19,7 +20,7 @@ def test_generate_event_data(): t_max = 1.0 event = nir.EventData(idx, time, n_neurons, t_max) node = nir.NIRNodeData({"spikes": event}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["spikes"].idx, idx) assert np.allclose(graph.nodes["node"].observables["spikes"].time, time) assert graph.nodes["node"].observables["spikes"].n_neurons == n_neurons @@ -34,7 +35,7 @@ def test_generate_valued_event_data(): t_max = 1.0 valued_event = nir.ValuedEventData(idx, time, n_neurons, t_max, value) node = nir.NIRNodeData({"current": valued_event}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["current"].idx, idx) assert np.allclose(graph.nodes["node"].observables["current"].time, time) assert np.allclose(graph.nodes["node"].observables["current"].value, value) @@ -70,7 +71,7 @@ def test_valued_conversion(): dt = 0.01 t_max = 0.3 valued_event = nir.ValuedEventData(idx, time, n_neurons, t_max, value) - gridded = valued_event.to_time_gridded(dt=dt) # noqa: F841 + gridded = valued_event.to_time_gridded(dt=dt) expected = np.zeros((1, 30, 2)) expected[0, 5, 0] = 1 expected[0, 10, 1] = 2 diff --git a/tests/test_ir.py b/tests/test_ir.py index e2653e1f..196c0eff 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -42,10 +42,6 @@ def test_eq(): b2 = nir.Input(np.array([2, 2])) o = nir.Output(np.array([2, 3])) - assert a == a - assert a2 == a2 - assert b == b - assert b2 == b2 assert a != a2 assert a != b assert a != b2 @@ -237,9 +233,7 @@ def test_flatten(): ir = nir.NIRGraph( nodes={ "in": nir.Input(input_type=np.array([4, 5, 2])), - "flat": nir.Flatten( - start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])} - ), + "flat": nir.Flatten(start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])}), "out": nir.Output(output_type=np.array([20, 2])), }, edges=[("in", "flat"), ("flat", "out")], @@ -252,51 +246,35 @@ def test_from_list_naming(): ir = nir.NIRGraph.from_list( nir.Linear(weight=np.array([[3, 1], [-1, 2], [1, 2]])), nir.Linear(weight=np.array([[3, 1], [-1, 4], [1, 2]]).T), - nir.Affine( - weight=np.array([[2, 1], [-1, 2], [1, 2]]), bias=np.array([1, 3, 2]) - ), - nir.Affine( - weight=np.array([[2, 1], [-1, 4], [1, 2]]).T, bias=np.array([-2, 2]) - ), + nir.Affine(weight=np.array([[2, 1], [-1, 2], [1, 2]]), bias=np.array([1, 3, 2])), + nir.Affine(weight=np.array([[2, 1], [-1, 4], [1, 2]]).T, bias=np.array([-2, 2])), nir.Linear(weight=np.array([[3, 1], [-1, 1], [1, 2]])), nir.Linear(weight=np.array([[3, 1], [-1, 3], [1, 2]]).T), - nir.Affine( - weight=np.array([[2, 1], [-1, 1], [1, 2]]), bias=np.array([1, 5, 2]) - ), - nir.Affine( - weight=np.array([[2, 1], [-1, 3], [1, 2]]).T, bias=np.array([-2, 3]) - ), + nir.Affine(weight=np.array([[2, 1], [-1, 1], [1, 2]]), bias=np.array([1, 5, 2])), + nir.Affine(weight=np.array([[2, 1], [-1, 3], [1, 2]]).T, bias=np.array([-2, 3])), ) - assert "input" in ir.nodes.keys() - assert "linear" in ir.nodes.keys() - assert "linear_1" in ir.nodes.keys() - assert "linear_2" in ir.nodes.keys() - assert "linear_3" in ir.nodes.keys() - assert "affine" in ir.nodes.keys() - assert "affine_1" in ir.nodes.keys() - assert "affine_2" in ir.nodes.keys() - assert "affine_3" in ir.nodes.keys() - assert "output" in ir.nodes.keys() + assert "input" in ir.nodes + assert "linear" in ir.nodes + assert "linear_1" in ir.nodes + assert "linear_2" in ir.nodes + assert "linear_3" in ir.nodes + assert "affine" in ir.nodes + assert "affine_1" in ir.nodes + assert "affine_2" in ir.nodes + assert "affine_3" in ir.nodes + assert "output" in ir.nodes assert np.allclose(ir.nodes["input"].input_type["input"], [2]) assert np.allclose(ir.nodes["linear"].weight, np.array([[3, 1], [-1, 2], [1, 2]])) - assert np.allclose( - ir.nodes["linear_1"].weight, np.array([[3, 1], [-1, 4], [1, 2]]).T - ) + assert np.allclose(ir.nodes["linear_1"].weight, np.array([[3, 1], [-1, 4], [1, 2]]).T) assert np.allclose(ir.nodes["affine"].weight, np.array([[2, 1], [-1, 2], [1, 2]])) assert np.allclose(ir.nodes["affine"].bias, np.array([1, 3, 2])) - assert np.allclose( - ir.nodes["affine_1"].weight, np.array([[2, 1], [-1, 4], [1, 2]]).T - ) + assert np.allclose(ir.nodes["affine_1"].weight, np.array([[2, 1], [-1, 4], [1, 2]]).T) assert np.allclose(ir.nodes["affine_1"].bias, np.array([-2, 2])) assert np.allclose(ir.nodes["linear_2"].weight, np.array([[3, 1], [-1, 1], [1, 2]])) - assert np.allclose( - ir.nodes["linear_3"].weight, np.array([[3, 1], [-1, 3], [1, 2]]).T - ) + assert np.allclose(ir.nodes["linear_3"].weight, np.array([[3, 1], [-1, 3], [1, 2]]).T) assert np.allclose(ir.nodes["affine_2"].weight, np.array([[2, 1], [-1, 1], [1, 2]])) assert np.allclose(ir.nodes["affine_2"].bias, np.array([1, 5, 2])) - assert np.allclose( - ir.nodes["affine_3"].weight, np.array([[2, 1], [-1, 3], [1, 2]]).T - ) + assert np.allclose(ir.nodes["affine_3"].weight, np.array([[2, 1], [-1, 3], [1, 2]]).T) assert np.allclose(ir.nodes["affine_3"].bias, np.array([-2, 3])) print(ir.nodes["output"].input_type["input"]) assert np.allclose(ir.nodes["output"].input_type["input"], [2]) @@ -387,9 +365,7 @@ def test_inputs_outputs_properties(): nodes={ "in1": nir.Input(np.array([4, 5, 2])), "in2": nir.Input(np.array([4, 5, 2])), - "flat": nir.Flatten( - start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])} - ), + "flat": nir.Flatten(start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])}), "out1": nir.Output(np.array([20, 2])), "out2": nir.Output(np.array([20, 2])), }, @@ -442,12 +418,8 @@ def test_sumpool_type_inference(): edges=[("input", "sumpool"), ("sumpool", "output")], ) assert np.array_equal(graph.output_type["output"], np.array([1, 32, 32])) - assert np.array_equal( - graph.nodes["output"].input_type["input"], np.array([1, 32, 32]) - ) - assert np.array_equal( - graph.nodes["output"].output_type["output"], np.array([1, 32, 32]) - ) + assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 32, 32])) + assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 32, 32])) def test_avgpool_type_inference(): @@ -464,12 +436,8 @@ def test_avgpool_type_inference(): edges=[("input", "avgpool"), ("avgpool", "output")], ) assert np.array_equal(graph.output_type["output"], np.array([1, 32, 32])) - assert np.array_equal( - graph.nodes["output"].input_type["input"], np.array([1, 32, 32]) - ) - assert np.array_equal( - graph.nodes["output"].output_type["output"], np.array([1, 32, 32]) - ) + assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 32, 32])) + assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 32, 32])) def test_flatten_type_inference(): @@ -514,18 +482,10 @@ def test_flatten_type_inference(): }, edges=[("input", "flatten"), ("flatten", "output")], ) - assert np.array_equal( - graph.nodes["flatten"].input_type["input"], test["input"] - ) - assert np.array_equal( - graph.nodes["flatten"].output_type["output"], test["output"] - ) - assert np.array_equal( - graph.nodes["output"].input_type["input"], test["output"] - ) - assert np.array_equal( - graph.nodes["output"].output_type["output"], test["output"] - ) + assert np.array_equal(graph.nodes["flatten"].input_type["input"], test["input"]) + assert np.array_equal(graph.nodes["flatten"].output_type["output"], test["output"]) + assert np.array_equal(graph.nodes["output"].input_type["input"], test["output"]) + assert np.array_equal(graph.nodes["output"].output_type["output"], test["output"]) assert np.array_equal(graph.input_type["input"], test["input"]) assert np.array_equal(graph.output_type["output"], test["output"]) @@ -594,15 +554,15 @@ def test_conv2d_type_inference(): except Exception as ex: raise AssertionError(f"type check failed for: {name}: {ex}") from ex - assert np.array_equal( - graph.nodes["output"].input_type["input"], np.array([1, 61, 61]) - ), name - assert np.array_equal( - graph.nodes["output"].output_type["output"], np.array([1, 61, 61]) - ), name - assert np.array_equal( - graph.nodes["conv"].output_type["output"], np.array([1, 61, 61]) - ), name + assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 61, 61])), ( + name + ) + assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 61, 61])), ( + name + ) + assert np.array_equal(graph.nodes["conv"].output_type["output"], np.array([1, 61, 61])), ( + name + ) assert np.array_equal(graph.input_type["input"], np.array([1, 64, 64])), name assert np.array_equal(graph.output_type["output"], np.array([1, 61, 61])), name @@ -654,15 +614,9 @@ def test_conv1d_type_inference(): except Exception as ex: raise AssertionError(f"type check failed for: {name}: {ex}") from ex - assert np.array_equal( - graph.nodes["output"].input_type["input"], np.array([1, 61]) - ), name - assert np.array_equal( - graph.nodes["output"].output_type["output"], np.array([1, 61]) - ), name - assert np.array_equal( - graph.nodes["conv"].output_type["output"], np.array([1, 61]) - ), name + assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 61])), name + assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 61])), name + assert np.array_equal(graph.nodes["conv"].output_type["output"], np.array([1, 61])), name assert np.array_equal(graph.input_type["input"], np.array([1, 64])), name assert np.array_equal(graph.output_type["output"], np.array([1, 61])), name @@ -697,18 +651,18 @@ def test_graph_input_output_type_inference(): assert ( graph.input_type is not None and len(graph.input_type) == 1 - and np.array_equal(list(graph.input_type.values())[0], np.array([6])) + and np.array_equal(next(iter(graph.input_type.values())), np.array([6])) ), f"unexpected graph input type for {name} after type inference" # Graph output should be set to the output_type of the Output node assert ( graph.output_type is not None and len(graph.output_type) == 1 - and np.array_equal(list(graph.output_type.values())[0], np.array([4])) + and np.array_equal(next(iter(graph.output_type.values())), np.array([4])) ), f"unexpected graph output type for {name} after type inference" # Input nodes should have input and output types set to the same values. assert len(graph.inputs) == 1, f"unexpected number of input nodes for {name}" - input_node = list(graph.inputs.values())[0] + input_node = next(iter(graph.inputs.values())) assert ( input_node.input_type is not None and len(input_node.input_type) == 1 @@ -724,7 +678,7 @@ def test_graph_input_output_type_inference(): # Output nodes should have input and output types set to the same values. assert len(graph.outputs) == 1, f"unexpected number of output nodes for {name}" - output_node = list(graph.outputs.values())[0] + output_node = next(iter(graph.outputs.values())) assert ( output_node.input_type is not None and len(output_node.input_type) == 1 @@ -745,9 +699,7 @@ def test_type_check_recurrent(): nodes={ "a": nir.Input(np.array([2])), "b": nir.Linear(np.random.rand(2, 2)), - "c": nir.IF( - r=np.random.rand(2), v_threshold=np.random.rand(2), v_reset=np.zeros(2) - ), + "c": nir.IF(r=np.random.rand(2), v_threshold=np.random.rand(2), v_reset=np.zeros(2)), "d": nir.Output(np.array([2])), }, edges=[("a", "b"), ("b", "c"), ("c", "b"), ("c", "d")], @@ -845,8 +797,6 @@ def test_validate_structure_recurrent_valid(): def test_node(): try: node = nir.ir.NIRNode() - assert ( - node is None - ), f"test failed, we should not be able to construct an NIRNode: {node}" + assert node is None, f"test failed, we should not be able to construct an NIRNode: {node}" except AttributeError: pass diff --git a/tests/test_readwrite.py b/tests/test_readwrite.py index 3e309273..5d327f92 100644 --- a/tests/test_readwrite.py +++ b/tests/test_readwrite.py @@ -23,11 +23,7 @@ def assert_equivalence(ir: nir.NIRGraph, ir2: nir.NIRGraph): assert_equivalence(ir.nodes[ik], ir2.nodes[ik]) else: for k, v in ir.nodes[ik].__dict__.items(): - if ( - isinstance(v, np.ndarray) - or isinstance(v, list) - or isinstance(v, tuple) - ): + if isinstance(v, (np.ndarray, list, tuple)): assert np.array_equal(v, getattr(ir2.nodes[ik], k)) elif isinstance(v, dict): d = getattr(ir2.nodes[ik], k) @@ -312,7 +308,7 @@ def test_deserialize(): nir.read(os.path.join(nir_base, file)) except Exception as e: print(f"Failed to read {file}: {e}") - raise e + raise def test_read_without_type_check(): diff --git a/tests/test_utils.py b/tests/test_utils.py index a6c5692e..8c64eff8 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,10 +1,10 @@ +import importlib + import numpy as np import pytest from nir.ir.utils import _index_tuple -import importlib - _TORCH_SPEC = importlib.util.find_spec("torch") is not None @@ -12,9 +12,7 @@ def test_index_tuple(): assert _index_tuple(1, 0) == 1 assert _index_tuple([1, 2], 0) == 1 assert _index_tuple(np.array([1, 2]), 0) == 1 - assert np.all( - np.equal(_index_tuple(np.array([[1, 2], [3, 4]]), 1), np.array([3, 4])) - ) + assert np.all(np.equal(_index_tuple(np.array([[1, 2], [3, 4]]), 1), np.array([3, 4]))) @pytest.mark.skipif(_TORCH_SPEC is not None, reason="requires torch") From 5468ef848426bc853d87c2fe497e8447693c83be Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Fri, 14 Aug 2026 17:38:15 +0200 Subject: [PATCH 2/4] chore: Update supported primitives documentation --- docs/source/supported_primitives.md | 25 +++++++++---------- docs/supported_primitives.py | 16 +++++++++--- .../examples/jaxsnn/supported_primitives.md | 5 ++-- .../examples/norse/supported_primitives.md | 4 +++ .../tmp/examples/spyx/supported_primitives.md | 10 +++----- 5 files changed, 36 insertions(+), 24 deletions(-) diff --git a/docs/source/supported_primitives.md b/docs/source/supported_primitives.md index 143a9227..85514450 100644 --- a/docs/source/supported_primitives.md +++ b/docs/source/supported_primitives.md @@ -3,26 +3,25 @@ This document lists which primitives are supported by the software frameworks fo - `→`: Supported for conversion from NIR - `←`: Supported for conversion to NIR - `⟷`: Supported for both conversion directions (to and from NIR) - Please note that this list is generated automatically and may not be entirely accurate.
| Primitive | hxtorch | jaxsnn | Lava | Nengo | Norse | rockpool | sinabs | snntorch | SpiNNaker2 | Spyx | |-----------|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:| -| Conv1d | | | | | | | ⟷ | | → | → | +| Conv1d | | | | | ← | | ⟷ | | → | | | Conv2d | | | → | | ⟷ | | ⟷ | ⟷ | → | ⟷ | -| Delay | | | | | | | | | | → | -| Flatten | | | → | | → | | ⟷ | ⟷ | → | → | -| Affine | | → | → | ⟷ | ⟷ | ⟷ | ⟷ | ⟷ | → | ⟷ | -| Linear | ⟷ | → | → | | ← | ⟷ | | ⟷ | → | ⟷ | -| Scale | | | | | | | | | | → | -| CubaLI | ⟷ | | | | | | | | | | -| CubaLIF | ⟷ | → | → | | ⟷ | ⟷ | | ⟷ | → | ⟷ | -| I | | | | | | | | | | → | +| Delay | | | | | | | | | | | +| Flatten | | | → | | ⟷ | | ⟷ | ⟷ | → | ⟷ | +| Affine | | | → | ⟷ | ⟷ | ⟷ | ⟷ | ⟷ | → | ⟷ | +| Linear | ⟷ | ⟷ | → | | ← | ⟷ | | ⟷ | → | ⟷ | +| Scale | | | | | | | | | | | +| CubaLI | ⟷ | | | | ← | | | | | | +| CubaLIF | ⟷ | ⟷ | → | | ⟷ | ⟷ | | ⟷ | → | ⟷ | +| I | | | | | | | | | | | | IF | | | → | | ⟷ | | ⟷ | → | → | ⟷ | | LI | | | | ⟷ | ⟷ | ⟷ | ⟷ | | | ← | | LIF | | | → | ⟷ | ⟷ | ⟷ | ⟷ | ⟷ | → | ⟷ | -| AvgPool2d | | | | | | | | ⟷ | | | -| SumPool2d | | | → | | → | | ⟷ | | → | → | -| Threshold | | | | | | | | | | → | \ No newline at end of file +| AvgPool2d | | | | | ← | | | ⟷ | | | +| SumPool2d | | | → | | → | | ⟷ | | → | ⟷ | +| Threshold | | | | | | | | | | ← | \ No newline at end of file diff --git a/docs/supported_primitives.py b/docs/supported_primitives.py index bdffe6f7..14b2ef26 100644 --- a/docs/supported_primitives.py +++ b/docs/supported_primitives.py @@ -25,7 +25,13 @@ ( "jaxsnn", "from_nir", - "https://raw.githubusercontent.com/electronicvisions/jaxsnn/refs/heads/main/src/pyjaxsnn/jaxsnn/event/from_nir.py", + "https://raw.githubusercontent.com/electronicvisions/jaxsnn/refs/heads/main/src/pyjaxsnn/jaxsnn/event/utils/from_nir.py", + None, + ), + ( + "jaxsnn", + "to_nir", + "https://raw.githubusercontent.com/electronicvisions/jaxsnn/refs/heads/main/src/pyjaxsnn/jaxsnn/event/utils/to_nir.py", None, ), ( @@ -104,7 +110,7 @@ "Spyx", "from_nir", "https://raw.githubusercontent.com/kmheckel/spyx/refs/heads/main/src/spyx/nir.py", - "_nir_node_to_spyx_node", + "_nir_node_to_spyx_module", ), ( "Spyx", @@ -124,7 +130,10 @@ response = requests.get(url) pattern = rf"def {function}\s*\(.*?\):(.*?)(?=\ndef |\Z)" match = re.search(pattern, response.text, re.DOTALL) - converter_contents[key] = match.group(0) + if match == None: + raise ValueError(f"Function {function} not found in {url}") + else: + converter_contents[key] = match.group(0) # Check which primitives are supported in each library/direction supported = {} @@ -171,6 +180,7 @@ - `←`: Supported for conversion to NIR - `⟷`: Supported for both conversion directions (to and from NIR) Please note that this list is generated automatically and may not be entirely accurate. +
""" full_md = static_md + "\n\n" + dynamic_md diff --git a/docs/tmp/examples/jaxsnn/supported_primitives.md b/docs/tmp/examples/jaxsnn/supported_primitives.md index 2ea915a0..07868c4e 100644 --- a/docs/tmp/examples/jaxsnn/supported_primitives.md +++ b/docs/tmp/examples/jaxsnn/supported_primitives.md @@ -1,8 +1,9 @@ ### Supported Primitives in jaxsnn -This library does not support conversion of any nodes to NIR. +This library supports conversion of the following nodes to NIR: +- Linear +- CubaLIF This library supports conversion of the following nodes from NIR: -- Affine - Linear - CubaLIF \ No newline at end of file diff --git a/docs/tmp/examples/norse/supported_primitives.md b/docs/tmp/examples/norse/supported_primitives.md index f088ae4a..5d3030fa 100644 --- a/docs/tmp/examples/norse/supported_primitives.md +++ b/docs/tmp/examples/norse/supported_primitives.md @@ -1,13 +1,17 @@ ### Supported Primitives in Norse This library supports conversion of the following nodes to NIR: +- Conv1d - Conv2d +- Flatten - Affine - Linear +- CubaLI - CubaLIF - IF - LI - LIF +- AvgPool2d This library supports conversion of the following nodes from NIR: - Conv2d diff --git a/docs/tmp/examples/spyx/supported_primitives.md b/docs/tmp/examples/spyx/supported_primitives.md index 819fe6a4..f1e79b8a 100644 --- a/docs/tmp/examples/spyx/supported_primitives.md +++ b/docs/tmp/examples/spyx/supported_primitives.md @@ -2,24 +2,22 @@ This library supports conversion of the following nodes to NIR: - Conv2d +- Flatten - Affine - Linear - CubaLIF - IF - LI - LIF +- SumPool2d +- Threshold This library supports conversion of the following nodes from NIR: -- Conv1d - Conv2d -- Delay - Flatten - Affine - Linear -- Scale - CubaLIF -- I - IF - LIF -- SumPool2d -- Threshold \ No newline at end of file +- SumPool2d \ No newline at end of file From c20fe686c8459326e96f5d9f616d331a6198ac7d Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Fri, 14 Aug 2026 17:44:08 +0200 Subject: [PATCH 3/4] chore: Adapt build.yml to new ruff --- .github/workflows/build.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 320ef275..d67231da 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -27,10 +27,8 @@ jobs: if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - name: Lint with ruff run: | - # stop the build if there are Python syntax errors or undefined names - ruff check . --output-format=github --select=E9,F63,F7,F82 --target-version=py37 # default set of ruff rules with GitHub Annotations - ruff check . --output-format=github --target-version=py37 --exclude=docs/ --exclude=paper/ + ruff check --output-format=github --target-version=py310 --exclude=docs/ --exclude=paper/ - name: Check formatting with black run: | black --check nir/ tests/ From 76439f1273f30100c06b6a6a22ee0a1d97ba5b15 Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Fri, 14 Aug 2026 17:49:01 +0200 Subject: [PATCH 4/4] chore: Make Black happy --- nir/data_ir/graph.py | 12 +++-- nir/ir/conv.py | 24 +++++++--- nir/ir/flatten.py | 4 +- nir/ir/graph.py | 68 ++++++++++++++++++-------- nir/ir/linear.py | 12 +++-- nir/ir/neuron.py | 21 +++++---- nir/ir/typing.py | 1 - nir/serialization.py | 28 ++++++++--- tests/test_ir.py | 110 ++++++++++++++++++++++++++++++------------- tests/test_utils.py | 4 +- 10 files changed, 202 insertions(+), 82 deletions(-) diff --git a/nir/data_ir/graph.py b/nir/data_ir/graph.py index 653735ff..873f3593 100644 --- a/nir/data_ir/graph.py +++ b/nir/data_ir/graph.py @@ -136,7 +136,9 @@ def to_time_gridded( Time step size. """ n_time_steps = int(self.t_max / dt) - discrete_data = np.zeros((self.n_samples, n_time_steps, self.n_neurons), dtype=bool) + discrete_data = np.zeros( + (self.n_samples, n_time_steps, self.n_neurons), dtype=bool + ) for sample in range(self.n_samples): valid_spikes = self.idx[sample] != -1 @@ -217,7 +219,9 @@ class NIRNodeData: def __post_init__(self): if not isinstance(self.observables, dict): - raise TypeError("observables must be a dictionary of EventData or TimeGriddedData") + raise TypeError( + "observables must be a dictionary of EventData or TimeGriddedData" + ) def __getitem__(self, idx): return self.observables[idx] @@ -274,4 +278,6 @@ def check_nodes(self, graph: NIRGraph): if not isinstance(graph_node, NIRNode): raise TypeError(f"Node {key} is not a NIRNode in the NIRGraph") if not node.check_observables(graph_node): - raise ValueError(f"Observables for node {key} do not match the NIRNode") + raise ValueError( + f"Observables for node {key} do not match the NIRNode" + ) diff --git a/nir/ir/conv.py b/nir/ir/conv.py index 8b93d425..dcb3e00d 100644 --- a/nir/ir/conv.py +++ b/nir/ir/conv.py @@ -43,14 +43,18 @@ class Conv1d(NIRNode): def __post_init__(self): if isinstance(self.padding, str) and self.padding not in ["same", "valid"]: - raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}") + raise ValueError( + f"padding must be 'same', 'valid', or int, not {self.padding}" + ) if self.input_shape is None: # leave input and output types undefined self.input_type = {"input": None} self.output_type = {"output": None} else: # infer input and output types from input_shape - self.input_type = {"input": np.array([self.weight.shape[1], self.input_shape])} + self.input_type = { + "input": np.array([self.weight.shape[1], self.input_shape]) + } output_shape = calculate_conv_output( self.input_shape, self.padding, @@ -58,7 +62,9 @@ def __post_init__(self): self.weight.shape[2], self.stride, ) - self.output_type = {"output": np.array([self.weight.shape[0], *output_shape])} + self.output_type = { + "output": np.array([self.weight.shape[0], *output_shape]) + } @dataclass(eq=False) @@ -98,7 +104,9 @@ class Conv2d(NIRNode): def __post_init__(self): if isinstance(self.padding, str) and self.padding not in ["same", "valid"]: - raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}") + raise ValueError( + f"padding must be 'same', 'valid', or int, not {self.padding}" + ) if isinstance(self.padding, int): self.padding = (self.padding, self.padding) if isinstance(self.stride, int): @@ -111,7 +119,9 @@ def __post_init__(self): self.output_type = {"output": None} else: # infer input and output types from input_shape - self.input_type = {"input": np.array([self.weight.shape[1], *self.input_shape])} + self.input_type = { + "input": np.array([self.weight.shape[1], *self.input_shape]) + } output_shape = calculate_conv_output( self.input_shape, self.padding, @@ -119,4 +129,6 @@ def __post_init__(self): self.weight.shape[2], self.stride, ) - self.output_type = {"output": np.array([self.weight.shape[0], *output_shape])} + self.output_type = { + "output": np.array([self.weight.shape[0], *output_shape]) + } diff --git a/nir/ir/flatten.py b/nir/ir/flatten.py index 23c29402..a1b262c4 100644 --- a/nir/ir/flatten.py +++ b/nir/ir/flatten.py @@ -35,7 +35,9 @@ def __post_init__(self): } # make sure input and output shape are valid if np.prod(self.input_type["input"]) != np.prod(self.output_type["output"]): - raise ValueError("input and output shape must have same number of elements") + raise ValueError( + "input and output shape must have same number of elements" + ) def to_dict(self) -> dict[str, Any]: ret = super().to_dict() diff --git a/nir/ir/graph.py b/nir/ir/graph.py index 625d37a9..7f706667 100644 --- a/nir/ir/graph.py +++ b/nir/ir/graph.py @@ -57,11 +57,15 @@ def __init__( @property def inputs(self): - return {name: node for name, node in self.nodes.items() if isinstance(node, Input)} + return { + name: node for name, node in self.nodes.items() if isinstance(node, Input) + } @property def outputs(self): - return {name: node for name, node in self.nodes.items() if isinstance(node, Output)} + return { + name: node for name, node in self.nodes.items() if isinstance(node, Output) + } @staticmethod def from_list(*nodes: NIRNode, type_check: bool = True) -> "NIRGraph": @@ -110,15 +114,23 @@ def __post_init__(self): self.metadata = {} def _update_input_output_types(self): - input_node_keys = [k for k, node in self.nodes.items() if isinstance(node, Input)] + input_node_keys = [ + k for k, node in self.nodes.items() if isinstance(node, Input) + ] self.input_type = ( - {node_key: self.nodes[node_key].input_type["input"] for node_key in input_node_keys} + { + node_key: self.nodes[node_key].input_type["input"] + for node_key in input_node_keys + } if len(input_node_keys) > 0 else None ) - output_node_keys = [k for k, node in self.nodes.items() if isinstance(node, Output)] + output_node_keys = [ + k for k, node in self.nodes.items() if isinstance(node, Output) + ] self.output_type = { - node_key: self.nodes[node_key].output_type["output"] for node_key in output_node_keys + node_key: self.nodes[node_key].output_type["output"] + for node_key in output_node_keys } def to_dict(self) -> dict[str, Any]: @@ -137,16 +149,20 @@ def from_dict(cls, kwargs: dict[str, Any]) -> "NIRGraph": assert "edges" in kwargs, "The incoming dictionary must hade a 'edges' entry" # Assert that the type is well-formed if "type" in kwargs: - assert kwargs["type"] == "NIRGraph", ( - "You are calling NIRGraph.from_dict with a different type " - ) + assert ( + kwargs["type"] == "NIRGraph" + ), "You are calling NIRGraph.from_dict with a different type " f"{type}. Either remove the entry or use .from_dict, such as Input.from_dict" kwargs_local["type"] = "NIRGraph" - kwargs_local["nodes"] = {k: dict2NIRNode(n) for k, n in kwargs_local["nodes"].items()} + kwargs_local["nodes"] = { + k: dict2NIRNode(n) for k, n in kwargs_local["nodes"].items() + } # h5py deserializes edges into a numpy array of type bytes and dtype=object, # hence using ensure_str here - kwargs_local["edges"] = [(ensure_str(a), ensure_str(b)) for a, b in kwargs_local["edges"]] + kwargs_local["edges"] = [ + (ensure_str(a), ensure_str(b)) for a, b in kwargs_local["edges"] + ] return super().from_dict(kwargs_local) def validate_structure(self): @@ -229,7 +245,9 @@ def check_types(self): post_repr = f"{edge[1]}.input: {post_input_type}" raise ValueError(f"type mismatch: {pre_repr} -> {post_repr}") else: - raise NotImplementedError("multiple input/output types not supported yet") + raise NotImplementedError( + "multiple input/output types not supported yet" + ) return True def infer_types(self): @@ -330,18 +348,26 @@ def infer_types(self): if undef_post_input_type: # define post input_type to be the same as pre output_type post_node.input_type = { - k.replace("output", "input"): v for k, v in pre_node.output_type.items() + k.replace("output", "input"): v + for k, v in pre_node.output_type.items() } elif type_mismatch: # set post input_type to be the same as pre output_type - pre_repr = f"{pre_key}.output: {np.array(list(pre_node.output_type.values()))}" - post_repr = f"{post_key}.input: {np.array(list(post_node.input_type.values()))}" - raise ValueError(f"Type inference error: type mismatch: {pre_repr} -> {post_repr}") + pre_repr = ( + f"{pre_key}.output: {np.array(list(pre_node.output_type.values()))}" + ) + post_repr = ( + f"{post_key}.input: {np.array(list(post_node.input_type.values()))}" + ) + raise ValueError( + f"Type inference error: type mismatch: {pre_repr} -> {post_repr}" + ) # make sure that output nodes have output_type = input_type if isinstance(post_node, Output): post_node.output_type = { - k.replace("input", "output"): v for k, v in post_node.input_type.items() + k.replace("input", "output"): v + for k, v in post_node.input_type.items() } # check if post output_type needs to be defined @@ -373,7 +399,9 @@ def infer_types(self): post_node.kernel_size, post_node.stride, ) - output_type = np.array([post_node.input_type["input"][0], *output_shape]) + output_type = np.array( + [post_node.input_type["input"][0], *output_shape] + ) post_node.output_type = {"output": output_type} elif isinstance(post_node, Flatten): @@ -386,7 +414,9 @@ def infer_types(self): } n_inputs = np.prod(post_node.input_type["input"]) n_outputs = np.prod(post_node.output_type["output"]) - assert n_inputs == n_outputs, "Flatten must not change the number of elements" + assert ( + n_inputs == n_outputs + ), "Flatten must not change the number of elements" seen.add(post_key) ready += [e for e in self.edges if e[0] == post_key and e[1] not in seen] diff --git a/nir/ir/linear.py b/nir/ir/linear.py index ffcf50af..2b4ea4b0 100644 --- a/nir/ir/linear.py +++ b/nir/ir/linear.py @@ -24,9 +24,13 @@ class Affine(NIRNode): def __post_init__(self): assert len(self.weight.shape) >= 2, "Weight must be at least 2D" self.input_type = { - "input": np.array(self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T)) + "input": np.array( + self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T) + ) + } + self.output_type = { + "output": np.array(self.weight.shape[:-2] + (self.weight.shape[-2],)) } - self.output_type = {"output": np.array(self.weight.shape[:-2] + (self.weight.shape[-2],))} @dataclass(eq=False) @@ -42,7 +46,9 @@ class Linear(NIRNode): def __post_init__(self): assert len(self.weight.shape) >= 2, "Weight must be at least 2D" self.input_type = { - "input": np.array(self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T)) + "input": np.array( + self.weight.shape[:-2] + tuple(np.array(self.weight.shape[-1:]).T) + ) } self.output_type = {"output": self.weight.shape[:-2] + (self.weight.shape[-2],)} diff --git a/nir/ir/neuron.py b/nir/ir/neuron.py index 573f332d..0a5a4a43 100644 --- a/nir/ir/neuron.py +++ b/nir/ir/neuron.py @@ -34,9 +34,12 @@ class CubaLI(NIRNode): w_in: np.ndarray = 1.0 # Input current weight def __post_init__(self): - assert self.tau_syn.shape == self.tau_mem.shape == self.r.shape == self.v_leak.shape, ( - "All parameters must have the same shape" - ) + assert ( + self.tau_syn.shape + == self.tau_mem.shape + == self.r.shape + == self.v_leak.shape + ), "All parameters must have the same shape" # If w_in is a scalar, make it an array of same shape as v_leak self.w_in = np.ones_like(self.v_leak) * self.w_in self.input_type = {"input": np.array(self.v_leak.shape)} @@ -155,9 +158,9 @@ class IF(NIRNode): def __post_init__(self): if self.v_reset is None: self.v_reset = np.zeros_like(self.v_threshold) - assert self.r.shape == self.v_threshold.shape == self.v_reset.shape, ( - "All parameters must have the same shape" - ) + assert ( + self.r.shape == self.v_threshold.shape == self.v_reset.shape + ), "All parameters must have the same shape" self.input_type = {"input": np.array(self.r.shape)} self.output_type = {"output": np.array(self.r.shape)} @@ -187,9 +190,9 @@ class LI(NIRNode): v_leak: np.ndarray # Leak voltage def __post_init__(self): - assert self.tau.shape == self.r.shape == self.v_leak.shape, ( - "All parameters must have the same shape" - ) + assert ( + self.tau.shape == self.r.shape == self.v_leak.shape + ), "All parameters must have the same shape" self.input_type = {"input": np.array(self.r.shape)} self.output_type = {"output": np.array(self.r.shape)} diff --git a/nir/ir/typing.py b/nir/ir/typing.py index 9e870838..b639e4b4 100644 --- a/nir/ir/typing.py +++ b/nir/ir/typing.py @@ -1,4 +1,3 @@ - import numpy as np from .node import NIRNode diff --git a/nir/serialization.py b/nir/serialization.py index 28cc5473..24082b8c 100644 --- a/nir/serialization.py +++ b/nir/serialization.py @@ -22,7 +22,9 @@ def try_byte_to_str(a: bytes | Any) -> str | Any: def read_node(node: Any) -> nir.NIRNode: """Read a graph from a HDF5 file.""" if node["type"][()] == b"Affine": - return nir.Affine(weight=node["weight"][()], bias=node["bias"][()], **_read_metadata(node)) + return nir.Affine( + weight=node["weight"][()], bias=node["bias"][()], **_read_metadata(node) + ) elif node["type"][()] == b"Conv1d": return nir.Conv1d( input_shape=(node["input_shape"][()] if "input_shape" in node else None), @@ -64,7 +66,9 @@ def read_node(node: Any) -> nir.NIRNode: return nir.Flatten( start_dim=node["start_dim"][()], end_dim=node["end_dim"][()], - input_type={"input": node["input_type"][()] if "input_type" in node else None}, + input_type={ + "input": node["input_type"][()] if "input_type" in node else None + }, **_read_metadata(node), ) elif node["type"][()] == b"I": @@ -73,13 +77,17 @@ def read_node(node: Any) -> nir.NIRNode: return nir.IF( r=node["r"][()], v_reset=( - node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] + if "v_reset" in node + else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], **_read_metadata(node), ) elif node["type"][()] == b"Input": - return nir.Input(input_type={"input": node["shape"][()]}, **_read_metadata(node)) + return nir.Input( + input_type={"input": node["shape"][()]}, **_read_metadata(node) + ) elif node["type"][()] == b"LI": return nir.LI( tau=node["tau"][()], @@ -95,7 +103,9 @@ def read_node(node: Any) -> nir.NIRNode: r=node["r"][()], v_leak=node["v_leak"][()], v_reset=( - node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] + if "v_reset" in node + else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], **_read_metadata(node), @@ -116,7 +126,9 @@ def read_node(node: Any) -> nir.NIRNode: r=node["r"][()], v_leak=node["v_leak"][()], v_reset=( - node["v_reset"][()] if "v_reset" in node else np.zeros_like(node["v_threshold"][()]) + node["v_reset"][()] + if "v_reset" in node + else np.zeros_like(node["v_threshold"][()]) ), v_threshold=node["v_threshold"][()], w_in=node["w_in"][()], @@ -129,7 +141,9 @@ def read_node(node: Any) -> nir.NIRNode: **_read_metadata(node), ) elif node["type"][()] == b"Output": - return nir.Output(output_type={"output": node["shape"][()]}, **_read_metadata(node)) + return nir.Output( + output_type={"output": node["shape"][()]}, **_read_metadata(node) + ) elif node["type"][()] == b"Scale": return nir.Scale(scale=node["scale"][()], **_read_metadata(node)) elif node["type"][()] == b"Threshold": diff --git a/tests/test_ir.py b/tests/test_ir.py index 196c0eff..b766b7b6 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -233,7 +233,9 @@ def test_flatten(): ir = nir.NIRGraph( nodes={ "in": nir.Input(input_type=np.array([4, 5, 2])), - "flat": nir.Flatten(start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])}), + "flat": nir.Flatten( + start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])} + ), "out": nir.Output(output_type=np.array([20, 2])), }, edges=[("in", "flat"), ("flat", "out")], @@ -246,12 +248,20 @@ def test_from_list_naming(): ir = nir.NIRGraph.from_list( nir.Linear(weight=np.array([[3, 1], [-1, 2], [1, 2]])), nir.Linear(weight=np.array([[3, 1], [-1, 4], [1, 2]]).T), - nir.Affine(weight=np.array([[2, 1], [-1, 2], [1, 2]]), bias=np.array([1, 3, 2])), - nir.Affine(weight=np.array([[2, 1], [-1, 4], [1, 2]]).T, bias=np.array([-2, 2])), + nir.Affine( + weight=np.array([[2, 1], [-1, 2], [1, 2]]), bias=np.array([1, 3, 2]) + ), + nir.Affine( + weight=np.array([[2, 1], [-1, 4], [1, 2]]).T, bias=np.array([-2, 2]) + ), nir.Linear(weight=np.array([[3, 1], [-1, 1], [1, 2]])), nir.Linear(weight=np.array([[3, 1], [-1, 3], [1, 2]]).T), - nir.Affine(weight=np.array([[2, 1], [-1, 1], [1, 2]]), bias=np.array([1, 5, 2])), - nir.Affine(weight=np.array([[2, 1], [-1, 3], [1, 2]]).T, bias=np.array([-2, 3])), + nir.Affine( + weight=np.array([[2, 1], [-1, 1], [1, 2]]), bias=np.array([1, 5, 2]) + ), + nir.Affine( + weight=np.array([[2, 1], [-1, 3], [1, 2]]).T, bias=np.array([-2, 3]) + ), ) assert "input" in ir.nodes assert "linear" in ir.nodes @@ -265,16 +275,24 @@ def test_from_list_naming(): assert "output" in ir.nodes assert np.allclose(ir.nodes["input"].input_type["input"], [2]) assert np.allclose(ir.nodes["linear"].weight, np.array([[3, 1], [-1, 2], [1, 2]])) - assert np.allclose(ir.nodes["linear_1"].weight, np.array([[3, 1], [-1, 4], [1, 2]]).T) + assert np.allclose( + ir.nodes["linear_1"].weight, np.array([[3, 1], [-1, 4], [1, 2]]).T + ) assert np.allclose(ir.nodes["affine"].weight, np.array([[2, 1], [-1, 2], [1, 2]])) assert np.allclose(ir.nodes["affine"].bias, np.array([1, 3, 2])) - assert np.allclose(ir.nodes["affine_1"].weight, np.array([[2, 1], [-1, 4], [1, 2]]).T) + assert np.allclose( + ir.nodes["affine_1"].weight, np.array([[2, 1], [-1, 4], [1, 2]]).T + ) assert np.allclose(ir.nodes["affine_1"].bias, np.array([-2, 2])) assert np.allclose(ir.nodes["linear_2"].weight, np.array([[3, 1], [-1, 1], [1, 2]])) - assert np.allclose(ir.nodes["linear_3"].weight, np.array([[3, 1], [-1, 3], [1, 2]]).T) + assert np.allclose( + ir.nodes["linear_3"].weight, np.array([[3, 1], [-1, 3], [1, 2]]).T + ) assert np.allclose(ir.nodes["affine_2"].weight, np.array([[2, 1], [-1, 1], [1, 2]])) assert np.allclose(ir.nodes["affine_2"].bias, np.array([1, 5, 2])) - assert np.allclose(ir.nodes["affine_3"].weight, np.array([[2, 1], [-1, 3], [1, 2]]).T) + assert np.allclose( + ir.nodes["affine_3"].weight, np.array([[2, 1], [-1, 3], [1, 2]]).T + ) assert np.allclose(ir.nodes["affine_3"].bias, np.array([-2, 3])) print(ir.nodes["output"].input_type["input"]) assert np.allclose(ir.nodes["output"].input_type["input"], [2]) @@ -365,7 +383,9 @@ def test_inputs_outputs_properties(): nodes={ "in1": nir.Input(np.array([4, 5, 2])), "in2": nir.Input(np.array([4, 5, 2])), - "flat": nir.Flatten(start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])}), + "flat": nir.Flatten( + start_dim=0, end_dim=1, input_type={"input": np.array([4, 5, 2])} + ), "out1": nir.Output(np.array([20, 2])), "out2": nir.Output(np.array([20, 2])), }, @@ -418,8 +438,12 @@ def test_sumpool_type_inference(): edges=[("input", "sumpool"), ("sumpool", "output")], ) assert np.array_equal(graph.output_type["output"], np.array([1, 32, 32])) - assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 32, 32])) - assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 32, 32])) + assert np.array_equal( + graph.nodes["output"].input_type["input"], np.array([1, 32, 32]) + ) + assert np.array_equal( + graph.nodes["output"].output_type["output"], np.array([1, 32, 32]) + ) def test_avgpool_type_inference(): @@ -436,8 +460,12 @@ def test_avgpool_type_inference(): edges=[("input", "avgpool"), ("avgpool", "output")], ) assert np.array_equal(graph.output_type["output"], np.array([1, 32, 32])) - assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 32, 32])) - assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 32, 32])) + assert np.array_equal( + graph.nodes["output"].input_type["input"], np.array([1, 32, 32]) + ) + assert np.array_equal( + graph.nodes["output"].output_type["output"], np.array([1, 32, 32]) + ) def test_flatten_type_inference(): @@ -482,10 +510,18 @@ def test_flatten_type_inference(): }, edges=[("input", "flatten"), ("flatten", "output")], ) - assert np.array_equal(graph.nodes["flatten"].input_type["input"], test["input"]) - assert np.array_equal(graph.nodes["flatten"].output_type["output"], test["output"]) - assert np.array_equal(graph.nodes["output"].input_type["input"], test["output"]) - assert np.array_equal(graph.nodes["output"].output_type["output"], test["output"]) + assert np.array_equal( + graph.nodes["flatten"].input_type["input"], test["input"] + ) + assert np.array_equal( + graph.nodes["flatten"].output_type["output"], test["output"] + ) + assert np.array_equal( + graph.nodes["output"].input_type["input"], test["output"] + ) + assert np.array_equal( + graph.nodes["output"].output_type["output"], test["output"] + ) assert np.array_equal(graph.input_type["input"], test["input"]) assert np.array_equal(graph.output_type["output"], test["output"]) @@ -554,15 +590,15 @@ def test_conv2d_type_inference(): except Exception as ex: raise AssertionError(f"type check failed for: {name}: {ex}") from ex - assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 61, 61])), ( - name - ) - assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 61, 61])), ( - name - ) - assert np.array_equal(graph.nodes["conv"].output_type["output"], np.array([1, 61, 61])), ( - name - ) + assert np.array_equal( + graph.nodes["output"].input_type["input"], np.array([1, 61, 61]) + ), name + assert np.array_equal( + graph.nodes["output"].output_type["output"], np.array([1, 61, 61]) + ), name + assert np.array_equal( + graph.nodes["conv"].output_type["output"], np.array([1, 61, 61]) + ), name assert np.array_equal(graph.input_type["input"], np.array([1, 64, 64])), name assert np.array_equal(graph.output_type["output"], np.array([1, 61, 61])), name @@ -614,9 +650,15 @@ def test_conv1d_type_inference(): except Exception as ex: raise AssertionError(f"type check failed for: {name}: {ex}") from ex - assert np.array_equal(graph.nodes["output"].input_type["input"], np.array([1, 61])), name - assert np.array_equal(graph.nodes["output"].output_type["output"], np.array([1, 61])), name - assert np.array_equal(graph.nodes["conv"].output_type["output"], np.array([1, 61])), name + assert np.array_equal( + graph.nodes["output"].input_type["input"], np.array([1, 61]) + ), name + assert np.array_equal( + graph.nodes["output"].output_type["output"], np.array([1, 61]) + ), name + assert np.array_equal( + graph.nodes["conv"].output_type["output"], np.array([1, 61]) + ), name assert np.array_equal(graph.input_type["input"], np.array([1, 64])), name assert np.array_equal(graph.output_type["output"], np.array([1, 61])), name @@ -699,7 +741,9 @@ def test_type_check_recurrent(): nodes={ "a": nir.Input(np.array([2])), "b": nir.Linear(np.random.rand(2, 2)), - "c": nir.IF(r=np.random.rand(2), v_threshold=np.random.rand(2), v_reset=np.zeros(2)), + "c": nir.IF( + r=np.random.rand(2), v_threshold=np.random.rand(2), v_reset=np.zeros(2) + ), "d": nir.Output(np.array([2])), }, edges=[("a", "b"), ("b", "c"), ("c", "b"), ("c", "d")], @@ -797,6 +841,8 @@ def test_validate_structure_recurrent_valid(): def test_node(): try: node = nir.ir.NIRNode() - assert node is None, f"test failed, we should not be able to construct an NIRNode: {node}" + assert ( + node is None + ), f"test failed, we should not be able to construct an NIRNode: {node}" except AttributeError: pass diff --git a/tests/test_utils.py b/tests/test_utils.py index 8c64eff8..bfbdb8da 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -12,7 +12,9 @@ def test_index_tuple(): assert _index_tuple(1, 0) == 1 assert _index_tuple([1, 2], 0) == 1 assert _index_tuple(np.array([1, 2]), 0) == 1 - assert np.all(np.equal(_index_tuple(np.array([[1, 2], [3, 4]]), 1), np.array([3, 4]))) + assert np.all( + np.equal(_index_tuple(np.array([[1, 2], [3, 4]]), 1), np.array([3, 4])) + ) @pytest.mark.skipif(_TORCH_SPEC is not None, reason="requires torch")