diff --git a/src/lava/proc/s4d/models.py b/src/lava/proc/s4d/models.py index 73241b29a..520f445f1 100644 --- a/src/lava/proc/s4d/models.py +++ b/src/lava/proc/s4d/models.py @@ -23,6 +23,7 @@ class S4dModel(PyLoihiProcessModel): a_in = LavaPyType(PyInPort.VEC_DENSE, float) s_out = LavaPyType(PyOutPort.VEC_DENSE, float) s4_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3) + inp_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3) # S4 variables s4_state: np.ndarray = LavaPyType(np.ndarray, complex) diff --git a/src/lava/proc/s4d/process.py b/src/lava/proc/s4d/process.py index 18ddc1d43..a99d3e431 100644 --- a/src/lava/proc/s4d/process.py +++ b/src/lava/proc/s4d/process.py @@ -18,7 +18,8 @@ def __init__( b: float, c: float, s4_state: ty.Optional[int] = 0, - s4_exp: ty.Optional[int] = 0) -> None: + s4_exp: ty.Optional[int] = 0, + inp_exp: ty.Optional[int] = 0) -> None: """ Neuron process that implements S4D (described by Gu et al., 2022) dynamics. @@ -55,6 +56,10 @@ def __init__( Scaling exponent with base 2 for the S4 state variables. Note: This should only be used for nc models. Default is 0. + inp_exp: int + Bit precision of the input signal. + Note: This should only be used for nc models. + Default is 0. """ super().__init__(shape=shape, @@ -62,7 +67,8 @@ def __init__( b=b, c=c, s4_state=s4_state, - s4_exp=s4_exp) + s4_exp=s4_exp, + inp_exp=inp_exp) # Ports self.a_in = InPort(shape=shape) self.s_out = OutPort(shape=shape) @@ -71,8 +77,9 @@ def __init__( self.a = Var(shape=shape, init=a) self.b = Var(shape=shape, init=b) self.c = Var(shape=shape, init=c) - self.s4_state = Var(shape=shape, init=0) + self.s4_state = Var(shape=shape, init=s4_state) self.s4_exp = Var(shape=(1,), init=s4_exp) + self.inp_exp = Var(shape=(1,), init=inp_exp) @property def shape(self) -> ty.Tuple[int, ...]: diff --git a/tests/lava/proc/s4d/test_process.py b/tests/lava/proc/s4d/test_process.py index 1c33b2357..cd6b5f2d3 100644 --- a/tests/lava/proc/s4d/test_process.py +++ b/tests/lava/proc/s4d/test_process.py @@ -14,17 +14,20 @@ def test_init(self) -> None: """Tests instantiation of S4d""" shape = 10 s4_exp = 12 + inp_exp = 8 a = np.ones(shape) * 0.5 b = np.ones(shape) * 0.8 c = np.ones(shape) * 0.9 s4d = S4d(shape=(shape,), s4_exp=s4_exp, + inp_exp=inp_exp, a=a, b=b, c=c) self.assertEqual(s4d.shape, (shape,)) self.assertEqual(s4d.s4_exp.init, s4_exp) + self.assertEqual(s4d.inp_exp.init, inp_exp) np.testing.assert_array_equal(s4d.a.init, a) np.testing.assert_array_equal(s4d.b.init, b) np.testing.assert_array_equal(s4d.c.init, c) diff --git a/tests/lava/proc/s4d/utils.py b/tests/lava/proc/s4d/utils.py index 2debe72c2..a323da338 100644 --- a/tests/lava/proc/s4d/utils.py +++ b/tests/lava/proc/s4d/utils.py @@ -24,107 +24,6 @@ def get_coefficients( return s4d_A, s4d_B, s4d_C -def run_bit_acc_model( - inp: np.ndarray, - num_steps: int, - model_dim: int, - d_states: int, - a: np.ndarray, - b: np.ndarray, - c: np.ndarray, - s4d_exp: int, - is_real: bool) -> Tuple[np.ndarray]: - """ - Run original S4d model in fixed precision. - - This function simulates the behavior of a linear time-invariant system - with diagonalized state-space representation. (S4D) - The state-space equations are given by: - s4_state_{k+1} = A * s4_state_k + B * input_k - out_k = C * s4_state_k - - where: - - s4_state_k is the state vector at time step k, - - input_k is the input vector at time step k, - - out_k is the output vector at time step k, - - A is the diagonal state matrix, - - B is the diagonal input matrix, - - C is the diagonal output matrix. - - The function computes the next output step of the - system for the given input signal. - - The function computes the output of the system for the given input signal - over num_steps time steps. - - Parameters - ---------- - inp: np.ndarray - Input signal to the model. - num_steps: int - Number of time steps to simulate the model. - model_dim: int - Dimensionality of the model. - d_states: int - Number of model states. - a: np.ndarray - Diagonal elements of the state matrix of the S4D model. - b: np.ndarray - Diagonal elements of the input matrix of the S4D model. - c: np.ndarray - Diagonal elements of the output matrix of the S4D model. - s4d_exp: int - Bit precision of a, b, c and the s4_state. - is_real: bool - Whether a, b, c and the s4_state are complex or real valued. - - Returns - ------- - Tuple[np.ndarray] - Tuple containing the output of the fixed precision model simulation. - """ - - a = a[:model_dim * d_states] - b = b[:model_dim * d_states] - c = c[:model_dim * d_states] - out = np.zeros((model_dim * d_states, num_steps)) - expansion_weights = np.kron(np.eye(model_dim), np.ones(d_states)) - expanded_inp = np.matmul(expansion_weights.T, inp) - - if is_real: - a = (a * 2**s4d_exp).astype(int) - b = (b * 2**s4d_exp).astype(int) - c = (c * 2**s4d_exp).astype(int) - s4_state = np.zeros((model_dim * d_states,)).flatten() - for idx, data_in in enumerate(expanded_inp.T): - s4_state = (s4_state * a * 2**-s4d_exp).astype(int) - + (data_in * b * 2**-s4d_exp).astype(int) - out[:, idx] = (c * s4_state * 2**-s4d_exp).astype(int) * 2 - else: - s4_state_real = np.zeros((1, model_dim * d_states)).astype(int) - s4_state_imag = np.zeros((1, model_dim * d_states)).astype(int) - a_imag = (a.imag * 2**s4d_exp).astype(int) - a_real = (a.real * 2**s4d_exp).astype(int) - b_imag = (b.imag * 2**s4d_exp).astype(int) - b_real = (b.real * 2**s4d_exp).astype(int) - c_real = (c.real * 2**s4d_exp).astype(int) - c_imag = (c.imag * 2**s4d_exp).astype(int) - - for idx, data_in in enumerate(expanded_inp.T): - s4_state_real_copy = np.copy(s4_state_real) - s4_state_imag_copy = np.copy(s4_state_imag) - s4_state_real = (s4_state_real * a_real * 2**-s4d_exp).astype(int) - - (s4_state_imag_copy * a_imag * 2**-s4d_exp).astype(int) - + (data_in * b_real * 2**-s4d_exp).astype(int) - s4_state_imag = ((s4_state_imag * a_real) * 2**-s4d_exp).astype(int) - + ((s4_state_real_copy * a_imag) * 2**-s4d_exp).astype(int) - + ((data_in * b_imag) * 2**-s4d_exp).astype(int) - out_val = ((c_real * s4_state_real) * 2**-s4d_exp).astype(int) - - ((c_imag * s4_state_imag) * 2**-s4d_exp).astype(int) - out[:, idx] = 2 * out_val - return out - - def run_original_model( inp: np.ndarray, num_steps: int,