How to use the simdkalman.primitives function in simdkalman

def test_vectorized(self):

        mean = np.zeros((3,2,1))
        mean[0,...] = np.array([[[1],[10]]]) # state 1
        mean[1,...] = np.array([[[2],[20]]]) # state 2
        mean[2,...] = np.array([[[3],[30]]]) # state 3

        stack_mats = lambda arr: np.vstack([a[np.newaxis,...] for a in arr])

        covariance = stack_mats([np.eye(2)]*3)
        state_transition = stack_mats([np.eye(2)]*3)
        process_noise = stack_mats([np.eye(2)]*3)*0.1

        m1, P1 = primitives.predict(
            mean,
            covariance,
            state_transition,
            process_noise)

        self.assertMatrixEqual(m1, mean)
        self.assertSequenceEqual(P1.shape, (3,2,2))

        observation_model = stack_mats([np.ones((1,2))]*3)
        observation_noise = stack_mats([np.eye(1)*0.1]*3)

        measurement = np.array([[[2]], [[np.nan]], [[33]]])

        m, P = primitives.update_with_nan_check(
            m1,
            P1,

def test_ensure_matrix(self):
        self.assertMatrixEqual(primitives.ensure_matrix(3), np.eye(1)*3)
        self.assertMatrixEqual(primitives.ensure_matrix([[3.0]]), np.eye(1)*3)
        self.assertMatrixEqual(primitives.ensure_matrix(np.eye(1)), np.eye(1))
        self.assertMatrixEqual(primitives.ensure_matrix(np.eye(3)), np.eye(3))