How to use the simdkalman.primitives.predict function in simdkalman

def test_predict(self):

        mean = np.array([[1],[2],[3]])
        covariance = np.eye(3)*2

        state_transition = np.eye(3)*0.5
        process_noise = np.eye(3)*0.1

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

        self.assertSequenceEqual(m1.shape, (3,1))
        self.assertSequenceEqual(P1.shape, (3,3))

        # should be diagonal
        self.assertTrue(np.linalg.norm(P1 - np.diag(np.diag(P1))) < 1e-6)

def test_one_dimensional(self):
        mean = np.array([[1]])
        covariance = np.array([[1]])
        state_transition = np.array([[2]])
        process_noise = np.array([[0.1]])

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

        self.assertMatrixEqual(m1, mean*2, epsilon=1e-6)
        self.assertSequenceEqual(P1.shape, (1,1))

        observation_model = np.array([[1]])
        observation_noise = np.array([[0.2]])

        measurement = np.array([[1]])

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

def test_semi_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 = np.eye(2)
        process_noise = np.eye(2)*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 = np.ones((1,2))
        observation_noise = np.eye(1)*0.1

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

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