jaxsnn.event.types

Classes

Any(*args, **kwargs)	Special type indicating an unconstrained type.
BaseModule(layer_idx)
BasePopulation()	Base class for all neuron populations across paradigms
BaseState()	Base class for all neuron states across paradigms
Carry(parameters, spikes, external_spikes, …)
Event(time, idx)
EventBaseModule(generator, …)	Base class for event-driven network modules.
EventData	alias of jaxsnn.event.types.Spike
EventT(*args, **kwargs)
Generic()	Abstract base class for generic types.
LayeredEventT(*args, **kwargs)
OptState(opt_state, params, …)	Container for the optimizer/training state across steps.
Population(generator, List[str], Dict[str, …)	Event-driven population module
PopulationFunctions(init, step, state, …)	Function container for event population modules
ProbeEvent(time, state, probe)
ProbeStepState(neuron_state, time, …)
Projection(generator, int], …)	Event-driven projection module
ProjectionFunctions(init, state, event, …)	Function container for event projection modules
Protocol()	Base class for protocol classes.
QueueHead	alias of jax.Array
QueueIndex	alias of jax.Array
SourcePopulation(generator, parameters, …)	Event-driven source population module
SourcePopulationFunctions(init, state, …)	Function container for event source population modules
Spike(time, idx, current, layer_idx, internal)
Step(parameters, spikes, external_spikes, …)
StepState(neuron_state, time)
TypeVar(name, *constraints[, bound, …])	Type variable.

Functions

jaxsnn.event.types.NamedTuple(typename, fields=None, /, **kwargs)

Typed version of namedtuple.

Usage:

class Employee(NamedTuple):
    name: str
    id: int

This is equivalent to:

Employee = collections.namedtuple('Employee', ['name', 'id'])

The resulting class has an extra __annotations__ attribute, giving a dict that maps field names to types. (The field names are also in the _fields attribute, which is part of the namedtuple API.) An alternative equivalent functional syntax is also accepted:

Employee = NamedTuple('Employee', [('name', str), ('id', int)])

jaxsnn.event.types.dataclass(cls=None, /, *, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False)

Add dunder methods based on the fields defined in the class.

Examines PEP 526 __annotations__ to determine fields.

If init is true, an __init__() method is added to the class. If repr is true, a __repr__() method is added. If order is true, rich comparison dunder methods are added. If unsafe_hash is true, a __hash__() method is added. If frozen is true, fields may not be assigned to after instance creation. If match_args is true, the __match_args__ tuple is added. If kw_only is true, then by default all fields are keyword-only. If slots is true, a new class with a __slots__ attribute is returned.

jaxsnn.event.types.field(*, default=<dataclasses._MISSING_TYPE object>, default_factory=<dataclasses._MISSING_TYPE object>, init=True, repr=True, hash=None, compare=True, metadata=None, kw_only=<dataclasses._MISSING_TYPE object>)

Return an object to identify dataclass fields.

default is the default value of the field. default_factory is a 0-argument function called to initialize a field’s value. If init is true, the field will be a parameter to the class’s __init__() function. If repr is true, the field will be included in the object’s repr(). If hash is true, the field will be included in the object’s hash(). If compare is true, the field will be used in comparison functions. metadata, if specified, must be a mapping which is stored but not otherwise examined by dataclass. If kw_only is true, the field will become a keyword-only parameter to __init__().

It is an error to specify both default and default_factory.

jaxsnn.event.types.tree_map(f: Callable[[…], Any], tree: Any, *rest: Any, is_leaf: Optional[Callable[[Any], bool], None] = None) → Any

Maps a multi-input function over pytree args to produce a new pytree.

Args:

f: function that takes 1 + len(rest) arguments, to be applied at the

corresponding leaves of the pytrees.

tree: a pytree to be mapped over, with each leaf providing the first

positional argument to f.

rest: a tuple of pytrees, each of which has the same structure as tree

or has tree as a prefix.

is_leaf: an optionally specified function that will be called at each

flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf.

Returns:

A new pytree with the same structure as tree but with the value at each leaf given by f(x, *xs) where x is the value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rest.

Examples:

>>> import jax.tree_util
>>> jax.tree_util.tree_map(lambda x: x + 1, {"x": 7, "y": 42})
{'x': 8, 'y': 43}

If multiple inputs are passed, the structure of the tree is taken from the first input; subsequent inputs need only have tree as a prefix:

>>> jax.tree_util.tree_map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]])
[[5, 7, 9], [6, 1, 2]]