Neurons¶

class ALIF(BaseNode):
    """Adaptive leaky integrate-and-fire neuron with conductance adaptation.

    The ALIF model extends standard LIF by adding a voltage-dependent
    potassium conductance (g_k) that creates spike-frequency adaptation.
    Each spike increases g_k by dg_k, which then decays exponentially.

    Dynamics:
        dv/dt = (-g_leak * (v - E_leak) - g_k * (v - E_k) + x) / c_m
        dg_k/dt = -g_k / tau_adapt

        At spike: g_k += dg_k

    Args:
        n_neuron: Number of neurons (int or tuple of dimensions).
        v_threshold: Firing threshold (mV). Default: 1.0.
        v_reset: Reset voltage after spike (mV). Default: 0.0.
        c_m: Membrane capacitance (pF). Default: 1.0.
        g_leak: Leak conductance (nS). Default: 1.0.
        E_leak: Leak reversal potential (mV). Default: 0.0.
        E_k: Potassium reversal potential (mV). Default: -70.0.
        g_k_init: Initial adaptation conductance (nS). Default: 0.0.
        tau_adapt: Adaptation time constant (ms). Default: 20.0.
        dg_k: Adaptation increment per spike (nS). Default: 0.0.
        tau_ref: Refractory period (ms). None disables refractory.
            Default: None.
        trainable_param: Set of parameter names to make trainable.
        surrogate_function: Surrogate gradient function. Default: Sigmoid().
        detach_reset: If True, detach reset signal. Default: False.
        hard_reset: If True, use hard reset. Default: False.
        pre_spike_v: If True, store pre-spike voltage. Default: False.
        step_mode: Step mode. Default: "s".
        backend: Backend implementation. Default: "torch".
        device: Device for tensors. Default: None.
        dtype: Data type for tensors. Default: None.

    Attributes:
        v: Membrane potential, shape (*batch, n_neuron).
        g_k: Adaptation conductance, shape (*batch, n_neuron).
        refractory: Refractory counter (if tau_ref specified).
        c_m, g_leak, E_leak, E_k: Neuron parameters.
        tau_adapt: Adaptation time constant.
        dg_k: Per-spike adaptation increment.
    """

    g_k: torch.Tensor
    refractory: torch.Tensor | None

    c_m: torch.Tensor | torch.nn.Parameter
    g_leak: torch.Tensor | torch.nn.Parameter
    E_leak: torch.Tensor | torch.nn.Parameter
    E_k: torch.Tensor | torch.nn.Parameter
    tau_adapt: torch.Tensor | torch.nn.Parameter
    dg_k: torch.Tensor | torch.nn.Parameter
    tau_ref: torch.Tensor | torch.nn.Parameter | None

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        v_threshold: float | Float[TensorLike, " n_neuron"] = 1.0,
        v_reset: float | Float[TensorLike, " n_neuron"] = 0.0,
        c_m: float | Float[TensorLike, " n_neuron"] = 1.0,
        g_leak: float | Float[TensorLike, " n_neuron"] = 1.0,
        E_leak: float | Float[TensorLike, " n_neuron"] = 0.0,
        E_k: float | Float[TensorLike, " n_neuron"] = -70.0,
        g_k_init: float | Float[TensorLike, " n_neuron"] = 0.0,
        tau_adapt: float | Float[TensorLike, " n_neuron"] = 20.0,
        dg_k: float | Float[TensorLike, " n_neuron"] = 0.0,
        tau_ref: float | Float[TensorLike, " n_neuron"] | None = None,
        trainable_param: set[str] = set(),
        surrogate_function: Callable = Sigmoid(),
        detach_reset: bool = False,
        hard_reset: bool = False,
        pre_spike_v: bool = False,
        step_mode: Literal["s"] = "s",
        backend: Literal["torch"] = "torch",
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__(
            n_neuron=n_neuron,
            v_threshold=v_threshold,
            v_reset=v_reset,
            trainable_param=trainable_param,
            surrogate_function=surrogate_function,
            detach_reset=detach_reset,
            hard_reset=hard_reset,
            pre_spike_v=pre_spike_v,
            step_mode=step_mode,
            backend=backend,
            device=device,
            dtype=dtype,
        )
        _factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
        self.def_param(
            "c_m",
            c_m,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "g_leak",
            g_leak,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "E_leak",
            E_leak,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "E_k",
            E_k,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "tau_adapt",
            tau_adapt,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "dg_k",
            dg_k,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self._use_refractory = tau_ref is not None
        if self._use_refractory:
            self.def_param(
                "tau_ref",
                tau_ref,
                trainable_param=self.trainable_param,
                **_factory_kwargs,
            )
            self.register_memory("refractory", 0.0, self.n_neuron)
        else:
            self.tau_ref = None

        self.register_memory("g_k", g_k_init, self.n_neuron)

    @property
    def v_rest(self):
        if self._v_rest is None:
            return self.v_reset
        return self._v_rest

    @v_rest.setter
    def v_rest(self, v_rest):
        if self._v_rest is not None:
            self._v_rest = v_rest

    @property
    def v_peak(self):
        return self.v_threshold

    @v_peak.setter
    def v_peak(self, value):
        self.v_threshold = value

    def dV(
        self,
        v: Float[Tensor, "*batch n_neuron"],
        g_k: Float[Tensor, "*batch n_neuron"],
        x: Float[Tensor, "*batch n_neuron"],
    ):
        leak_term = -self.g_leak * (v - self.E_leak)
        adapt_term = -g_k * (v - self.E_k)
        derivative = (leak_term + adapt_term + x) / self.c_m
        linear = (-self.g_leak - g_k) / self.c_m
        return derivative, linear

    def dgk(
        self, g_k: Float[Tensor, "*batch n_neuron"]
    ) -> Float[Tensor, "*batch n_neuron"]:
        derivative = -g_k / self.tau_adapt
        linear = -1.0 / self.tau_adapt
        return derivative, linear

    def neuronal_charge(self, x: Float[Tensor, "*batch n_neuron"]):
        dt = environ.get("dt")
        self.v = exp_euler_step(self.dV, self.v, self.g_k, x, dt=dt)

    def neuronal_adaptation(self):
        dt = environ.get("dt")
        self.g_k = exp_euler_step(self.dgk, self.g_k, dt=dt)

    def neuronal_fire(self):
        spike = self.surrogate_function(
            (self.v - self.v_threshold) / (self.v_threshold - self.v_reset)
        )
        if not self._use_refractory:
            return spike
        not_in_refractory = self.refractory == 0
        spike = spike * not_in_refractory.detach().to(self.v.dtype)
        return spike

    def neuronal_reset(self, spike: Float[Tensor, "*batch n"]):
        if self.detach_reset:
            spike_d = spike.detach()
        else:
            spike_d = spike

        if self.pre_spike_v:
            self.v_pre_spike = self.v.clone()

        if self.hard_reset:
            self.v = self.v - (self.v - self.v_reset) * spike_d
        else:
            self.v = self.v - (self.v_threshold - self.v_reset) * spike_d

        self.g_k = self.g_k + self.dg_k * spike_d

        if self._use_refractory:
            self.refractory = torch.relu(
                self.refractory + spike_d * self.tau_ref - environ.get("dt")
            )

    def extra_repr(self):
        g_k_init = self._memories_rv["g_k"].value
        parts = [
            f"c_m={self._format_repr_value(self.c_m)}",
            f"g_leak={self._format_repr_value(self.g_leak)}",
            f"E_leak={self._format_repr_value(self.E_leak)}",
            f"E_k={self._format_repr_value(self.E_k)}",
            f"tau_adapt={self._format_repr_value(self.tau_adapt)}",
            f"dg_k={self._format_repr_value(self.dg_k)}",
            f"g_k_init={self._format_repr_value(g_k_init)}",
            f"tau_ref={self._format_repr_value(self.tau_ref)}"
            if self._use_refractory
            else "tau_ref=None",
        ]
        base = super().extra_repr()
        if base:
            parts.append(base)
        return ", ".join(parts)

class ELIF(ALIF):
    """Exponential integrate-and-fire neuron with adaptation.

    The ELIF model extends ALIF by adding an exponential term to the
    voltage dynamics, creating a sharp upswing when approaching threshold
    (the "initiation zone"). This captures the rapid depolarization seen
    in real neurons.

    Dynamics:
        dv/dt = (g_leak * delta_T * exp((v - v_T) / delta_T)
                 - g_leak * (v - E_leak) - g_k * (v - E_k) + x) / c_m
        dg_k/dt = -g_k / tau_adapt

    The exponential term creates a soft threshold effect where membrane
    potential accelerates as it approaches v_T.

    Args:
        n_neuron: Number of neurons (int or tuple of dimensions).
        v_threshold: Firing threshold (mV). Default: 1.0.
        v_reset: Reset voltage after spike (mV). Default: 0.0.
        c_m: Membrane capacitance (pF). Default: 1.0.
        g_leak: Leak conductance (nS). Default: 1.0.
        E_leak: Leak reversal potential (mV). Default: 0.0.
        E_k: Potassium reversal potential (mV). Default: -70.0.
        g_k_init: Initial adaptation conductance (nS). Default: 0.0.
        tau_adapt: Adaptation time constant (ms). Default: 20.0.
        dg_k: Adaptation increment per spike (nS). Default: 0.0.
        tau_ref: Refractory period (ms). Default: 0.0.
        delta_T: Slope factor for exponential term (mV). Default: 1.0.
        v_T: Soft threshold potential (mV). Default: 0.0.
        trainable_param: Set of parameter names to make trainable.
        surrogate_function: Surrogate gradient function. Default: Sigmoid().
        detach_reset: If True, detach reset signal. Default: False.
        hard_reset: If True, use hard reset. Default: False.
        pre_spike_v: If True, store pre-spike voltage. Default: False.
        step_mode: Step mode. Default: "s".
        backend: Backend implementation. Default: "torch".
        device: Device for tensors. Default: None.
        dtype: Data type for tensors. Default: None.

    Attributes:
        delta_T: Slope factor for exponential term.
        v_T: Soft threshold potential.
    """

    delta_T: torch.Tensor | torch.nn.Parameter
    v_T: torch.Tensor | torch.nn.Parameter

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        v_threshold: float | Float[TensorLike, " n_neuron"] = 1.0,
        v_reset: float | Float[TensorLike, " n_neuron"] = 0.0,
        c_m: float | Float[TensorLike, " n_neuron"] = 1.0,
        g_leak: float | Float[TensorLike, " n_neuron"] = 1.0,
        E_leak: float | Float[TensorLike, " n_neuron"] = 0.0,
        E_k: float | Float[TensorLike, " n_neuron"] = -70.0,
        g_k_init: float | Float[TensorLike, " n_neuron"] = 0.0,
        tau_adapt: float | Float[TensorLike, " n_neuron"] = 20.0,
        dg_k: float | Float[TensorLike, " n_neuron"] = 0.0,
        tau_ref: float | Float[TensorLike, " n_neuron"] | None = 0.0,
        delta_T: float | Float[TensorLike, " n_neuron"] = 1.0,
        v_T: float | Float[TensorLike, " n_neuron"] = 0.0,
        trainable_param: set[str] = set(),
        surrogate_function: Callable = Sigmoid(),
        detach_reset: bool = False,
        hard_reset: bool = False,
        pre_spike_v: bool = False,
        step_mode: Literal["s"] = "s",
        backend: Literal["torch"] = "torch",
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__(
            n_neuron=n_neuron,
            v_threshold=v_threshold,
            v_reset=v_reset,
            c_m=c_m,
            g_leak=g_leak,
            E_leak=E_leak,
            E_k=E_k,
            g_k_init=g_k_init,
            tau_adapt=tau_adapt,
            dg_k=dg_k,
            tau_ref=tau_ref,
            trainable_param=trainable_param,
            surrogate_function=surrogate_function,
            detach_reset=detach_reset,
            hard_reset=hard_reset,
            pre_spike_v=pre_spike_v,
            step_mode=step_mode,
            backend=backend,
            device=device,
            dtype=dtype,
        )
        _factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
        self.def_param(
            "delta_T",
            delta_T,
            sizes=self.n_neuron,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "v_T",
            v_T,
            sizes=self.n_neuron,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )

    def dV(
        self,
        v: Float[Tensor, "*batch n_neuron"],
        g_k: Float[Tensor, "*batch n_neuron"],
        x: Float[Tensor, "*batch n_neuron"],
    ):
        leak_term = -self.g_leak * (v - self.E_leak)
        adapt_term = -g_k * (v - self.E_k)
        exp_term = self.g_leak * self.delta_T * torch.exp((v - self.v_T) / self.delta_T)
        derivative = (leak_term + adapt_term + exp_term + x) / self.c_m
        linear = (-self.g_leak - g_k + exp_term / self.delta_T) / self.c_m
        return derivative, linear

    def neuronal_charge(self, x: Float[Tensor, "*batch n_neuron"]):
        dt = environ.get("dt")
        self.v = exp_euler_step(self.dV, self.v, self.g_k, x, dt=dt)

    def extra_repr(self):
        parts = [
            f"delta_T={self._format_repr_value(self.delta_T)}",
            f"v_T={self._format_repr_value(self.v_T)}",
        ]
        base = super().extra_repr()
        if base:
            parts.append(base)
        return ", ".join(parts)

class GLIF3(BaseNode):
    """GLIF3 model with after-spike currents and refractory period.

    The GLIF3 model extends standard LIF by adding after-spike currents
    (ASC) that capture spike-frequency adaptation. Each spike adds
    asc_amps to the ASC vector, which then decays exponentially with
    time constants 1/k.

    Dynamics:
        dV/dt = -(V - V_rest) / tau + (I_in + sum(I_asc)) / c_m
        dI_asc/dt = -k * I_asc

        At spike: I_asc += asc_amps

    Args:
        n_neuron: Number of neurons (int or tuple of dimensions).
        v_threshold: Firing threshold (mV). Default: -50.0.
        v_reset: Reset voltage after spike (mV). Default: -70.0.
        v_rest: Resting potential (mV). Defaults to v_reset if None.
        c_m: Membrane capacitance (pF). Default: 0.05.
        tau: Membrane time constant (ms). Default: 20.0.
        k: ASC decay rates (ms^-1), can be list for multiple ASC components.
            Default: [0.2].
        asc_amps: ASC amplitudes (pA) added at each spike.
            Default: [0.0].
        tau_ref: Refractory period (ms). Default: 0.0.
        trainable_param: Set of parameter names to make trainable.
        surrogate_function: Surrogate gradient function.
            Default: ``Erf(alpha=4, damping_factor=0.5)``, matching the
            Gaussian surrogate used in Chen et al. (2022).
        detach_reset: If True, detach reset signal. Default: False.
        hard_reset: If True, use hard reset. Default: False.
        pre_spike_v: If True, store pre-spike voltage. Default: False.
        step_mode: Step mode. Default: "s".
        backend: Backend implementation. Default: "torch".
        device: Device for tensors. Default: None.
        dtype: Data type for tensors. Default: None.

    Attributes:
        v: Membrane potential, shape (*batch, n_neuron).
        Iasc: After-spike currents, shape (*batch, n_neuron, n_Iasc).
        refractory: Refractory counter (if tau_ref > 0).
        c_m, tau, tau_ref: Neuron parameters.
        k: ASC decay rates, shape (n_neuron, n_Iasc) or (n_Iasc,).
        asc_amps: ASC amplitudes, shape (n_neuron, n_Iasc) or (n_Iasc,).
        n_Iasc: Number of ASC components.

    References:
        Teeter et al., "Generalized leaky integrate-and-fire models classify
        multiple neuron types," *Nature Communications*, 2018.

        Chen, G., Scherr, F., & Maass, W. (2022). A data-based large-scale
        model for primary visual cortex enables brain-like robust and versatile
        visual processing. *Science Advances*, 8(44), eabq7592.
        https://doi.org/10.1126/sciadv.abq7592
    """

    # make mypy typing and autocompletion easier
    Iasc: torch.Tensor
    refractory: torch.Tensor | None

    c_m: torch.Tensor | torch.nn.Parameter
    tau: torch.Tensor | torch.nn.Parameter
    tau_ref: torch.Tensor | torch.nn.Parameter | None
    k: torch.Tensor | torch.nn.Parameter
    asc_amps: torch.Tensor | torch.nn.Parameter

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        v_threshold: float | Float[TensorLike, " n_neuron"] = -50.0,  # mV
        v_reset: float | Float[TensorLike, " n_neuron"] = -70.0,  # mV
        v_rest: None | float | Float[TensorLike, " n_neuron"] = None,
        c_m: float | Float[TensorLike, " n_neuron"] = 0.05,  # 1/20 pfarad
        tau: float | Float[TensorLike, " n_neuron"] = 20.0,  # ms
        k: float | Sequence[float] | Float[TensorLike, "n_neuron {self.n_Iasc}"] = [
            0.2
        ],  # ms^-1
        asc_amps: float
        | Sequence[float]
        | Float[TensorLike, "n_neuron {self.n_Iasc}"] = [0.0],  # pA
        tau_ref: float | Float[TensorLike, " n_neuron"] | None = 0.0,  # ms
        trainable_param: set[str] = set(),
        surrogate_function: Callable = Erf(alpha=4.0, damping_factor=0.5),
        detach_reset: bool = False,
        hard_reset: bool = False,
        pre_spike_v: bool = False,
        step_mode: Literal["s"] = "s",
        backend: Literal["torch"] = "torch",
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__(
            n_neuron=n_neuron,
            v_threshold=v_threshold,
            v_reset=v_reset,
            trainable_param=trainable_param,
            surrogate_function=surrogate_function,
            detach_reset=detach_reset,
            step_mode=step_mode,
            backend=backend,
            pre_spike_v=pre_spike_v,
        )
        _factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
        self.hard_reset = hard_reset
        self.def_param(
            "c_m",
            c_m,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self.def_param(
            "tau",
            tau,
            trainable_param=self.trainable_param,
            **_factory_kwargs,
        )
        self._use_refractory = tau_ref is not None
        if self._use_refractory:
            self.def_param(
                "tau_ref",
                tau_ref,
                trainable_param=self.trainable_param,
                **_factory_kwargs,
            )
            self.register_memory("refractory", 0.0, self.n_neuron)
        else:
            self.tau_ref = None

        # for compat
        if v_rest is not None:
            self.def_param(
                "_v_rest",
                v_rest,
                trainable_param=self.trainable_param,
                **_factory_kwargs,
            )
        else:
            self._v_rest = None

        # Handle after-spike currents.
        if isinstance(asc_amps, Number):
            asc_amps = [asc_amps]
        if isinstance(k, Number):
            k = [k]

        resolved_asc_sizes = self.def_param_resolve_sizes(
            k,
            asc_amps,
            sizes=self.n_neuron + (None,),
        )
        self.n_Iasc: int = resolved_asc_sizes[-1]

        self.def_param(
            "k",
            k,
            sizes=resolved_asc_sizes,
            trainable_param=self.trainable_param,
            normalize_to_sizes=True,
            **_factory_kwargs,
        )
        self.def_param(
            "asc_amps",
            asc_amps,
            sizes=resolved_asc_sizes,
            trainable_param=self.trainable_param,
            normalize_to_sizes=True,
            **_factory_kwargs,
        )

        self.register_memory(
            "Iasc",
            [
                0.0,
            ]
            * self.n_Iasc,
            self.n_neuron + (self.n_Iasc,),
        )

    @property
    def v_rest(self) -> torch.Tensor:
        """Resting potential (mV).

        For compatibility with GLIF4/GLIF5, falls back to v_reset if
        not explicitly set during initialization.

        Returns:
            Resting potential tensor.
        """
        if self._v_rest is None:
            return self.v_reset
        return self._v_rest

    @v_rest.setter
    def v_rest(self, v_rest: float | torch.Tensor):
        """Set resting potential.

        Args:
            v_rest: New resting potential value (mV).
        """
        if self._v_rest is not None:
            self._v_rest = v_rest

    def dIasc(self, Iasc: Float[Tensor, "*batch n_neuron {self.n_Iasc}"]) -> tuple:
        """Compute ASC derivative for exponential Euler integration.

        Args:
            Iasc: After-spike currents, shape (*batch, n_neuron, n_Iasc).

        Returns:
            Tuple of (derivative, linear_coefficient) for exp_euler_step.
        """
        return -self.k * Iasc, -self.k

    def dV(
        self,
        v: Float[Tensor, "*batch n_neuron"],
        Iasc: Float[Tensor, "*batch n_neuron {self.n_Iasc}"],
        x: Float[Tensor, "*batch n_neuron"],
    ) -> tuple:
        """Compute membrane potential derivative for exp Euler integration.

        Args:
            v: Membrane potential, shape (*batch, n_neuron).
            Iasc: After-spike currents, shape (*batch, n_neuron, n_Iasc).
            x: Input current, shape (*batch, n_neuron).

        Returns:
            Tuple of (derivative, linear_coefficient) for exp_euler_step.
        """
        Isum = x
        # torch.autocast will cast half to float32 for sum op
        # see https://docs.pytorch.org/docs/stable/amp.html#ops-that-can-autocast-to-float32
        # here Iasc generally only have <4 modes, so no overflow guaranteed
        return (
            -(v - self.v_rest) / self.tau
            + (Isum + Iasc.sum(-1, dtype=Iasc.dtype)) / self.c_m,
            -1.0 / self.tau,
        )

    def neuronal_charge(self, x: Float[Tensor, "*batch n_neuron"]):
        v = exp_euler_step(self.dV, self.v, self.Iasc, x, dt=environ.get("dt"))
        self.v = v

    def neuronal_adaptation(self):
        self.Iasc = exp_euler_step(self.dIasc, self.Iasc, dt=environ.get("dt"))

    def neuronal_fire(self):
        # Check if voltage exceeds threshold and not in refractory period
        spike = self.surrogate_function(
            (self.v - self.v_threshold) / (self.v_threshold - self.v_reset)
        )
        if not self._use_refractory:
            return spike
        not_in_refractory = self.refractory == 0
        spike = spike * not_in_refractory.detach().to(self.v.dtype)
        return spike

    def neuronal_reset(self, spike: Float[Tensor, "*batch n"]):
        if self.detach_reset:
            spike_d = spike.detach()
        else:
            spike_d = spike

        if self.pre_spike_v:
            self.v_pre_spike = self.v.clone()

        if self.hard_reset:
            # hard reset
            self.v = self.v - (self.v - self.v_reset) * spike_d
        else:
            # soft reset
            self.v = self.v - (self.v_threshold - self.v_reset) * spike_d

        # Add after-spike currents
        self.Iasc = self.Iasc + self.asc_amps * spike_d[..., None]

        if self._use_refractory:
            # Set refractory period
            self.refractory = torch.relu(
                self.refractory + spike_d * self.tau_ref - environ.get("dt")
            )

    def get_rheobase(self):
        """Calculate rheobase current, the minimum constant input current
        required to make the neuron fire."""
        return get_rheobase(self.v_threshold, self.v_rest, self.c_m, self.tau)

    def extra_repr(self):
        parts = [
            f"c_m={self._format_repr_value(self.c_m)}",
            f"tau={self._format_repr_value(self.tau)}",
            f"tau_ref={self._format_repr_value(self.tau_ref)}"
            if self._use_refractory
            else "tau_ref=None",
            f"n_Iasc={self.n_Iasc}",
            f"k={self._format_repr_value(self.k)}",
            f"asc_amps={self._format_repr_value(self.asc_amps)}",
            "v_rest=auto"
            if self._v_rest is None
            else f"v_rest={self._format_repr_value(self._v_rest)}",
        ]
        base = super().extra_repr()
        if base:
            parts.append(base)
        return ", ".join(parts)

    # TODO: headache to define precise input-output shapes
    # TODO: shape handling not torch.compile friendly
    def _normalize_state_shapes(
        self,
        x: TensorLike | float,
        v0: TensorLike | float,
        Iasc0: TensorLike | float,
        dt: TensorLike | float,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        device = device or self.v_reset.device
        dtype = dtype or self.v_reset.dtype
        x, v0, Iasc0 = (
            torch.as_tensor(x, device=device, dtype=dtype),
            torch.as_tensor(v0, device=device, dtype=dtype),
            torch.as_tensor(Iasc0, device=device, dtype=dtype),
        )
        if isinstance(dt, float):
            dt = torch.tensor([dt], device=device, dtype=dtype)
        else:
            dt = torch.as_tensor(dt, device=device, dtype=dtype)

        shapes = (x.shape, v0.shape, Iasc0.shape[:-1])
        longest_shape = max(shapes, key=len)
        if dt.shape[0] != longest_shape[0]:
            dt = expand_trailing_dims(dt, longest_shape, broadcast_only=True)

        return x, v0, Iasc0, dt

    def forward_exact_no_spike(
        self,
        x: Float[Tensor, "*batch #neuron"] | Float[Tensor, "*batch"],
        v0: Float[Tensor, "*batch neuron"] | None = None,
        Iasc0: Float[Tensor, "*batch neuron {self.n_Iasc}"] | None = None,
        dt: float
        | Float[TensorLike, "#time *batch neuron"]
        | Float[TensorLike, "#*batch neuron"]
        | Float[TensorLike, "#time *batch"]
        | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if dt is None:
            dt = environ.get("dt")

        update = (v0 is None) and (Iasc0 is None)
        if v0 is None:
            v0 = self.v
        if Iasc0 is None:
            Iasc0 = self.Iasc

        x, v0, Iasc0, dt = self._normalize_state_shapes(x, v0, Iasc0, dt)

        v_inf = self.v_reset + x * self.tau / self.c_m

        exp_m = torch.exp(-dt / self.tau)
        # (time, batch, neuron, n_Iasc)
        exp_asc = torch.exp(-dt[..., None] * self.k)

        Iasc = Iasc0 * exp_asc

        # degenerate case if tau=tau_asc=1/k
        Iasc_contrib = torch.where(
            torch.abs(self.k - 1 / self.tau[..., None]) > 1e-12,
            (Iasc0 / self.c_m[..., None])
            * (exp_asc - exp_m[..., None])
            / (1.0 / self.tau[..., None] - self.k),
            (Iasc0 / self.c_m[..., None]) * (dt * exp_m)[..., None],
        )
        v = v_inf + (v0 - v_inf) * exp_m + Iasc_contrib.sum(dim=-1)

        if update:
            self.v = v
            self.Iasc = Iasc
        return v, Iasc