jaxsnn.discrete.functional.lif

Classes

BaseState()

Base class for all neuron states across paradigms

DenseData

alias of jax.Array

LIFParameters(tau_syn, tau_mem, v_th, …)

LIFState(V, I)

Parameter

alias of jax.Array

Functions

jaxsnn.discrete.functional.lif.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.discrete.functional.lif.lif_step(inputs: Dict[str, jax.Array], state: jaxsnn.discrete.functional.lif.LIFState, parameters: Optional[jax.Array, None], method: Callable, v_leak: float, v_th: float, v_reset: float, tau_mem: float, tau_syn: float, dt: float = 0.001) → Tuple[jaxsnn.discrete.functional.lif.LIFState, jax.Array]

Euler step of a leaky-integrate-and-fire neuron.

Parameters

  • inputs – Dictionary of input currents

  • state – Current neuron state

  • parameters – Optional learnable parameters

  • method – Surrogate gradient method for the threshold function

  • v_leak – Leak potential

  • v_th – Threshold potential

  • v_reset – Reset potential

  • tau_mem – Membrane time constant

  • tau_syn – Synaptic time constant

  • dt – Time step size

Returns

Tuple of updated neuron state and membrane potential

jaxsnn.discrete.functional.lif.tree_reduce(function: Callable[[T, Any], T], tree: Any, *, is_leaf: Optional[Callable[[Any], bool], None] = 'None') → T
jaxsnn.discrete.functional.lif.tree_reduce(function: Callable[[T, Any], T], tree: Any, initializer: T, is_leaf: Optional[Callable[[Any], bool], None] = None) → T