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

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,
            observation_model,
            observation_noise,
            measurement)

        self.assertSequenceEqual(m.shape, (3,2,1))
        self.assertSequenceEqual(P.shape, (3,2,2))

        self.assertMatrixEqual(m[1,...], m1[1,...])
        self.assertMatrixEqual(m[2,...], mean[2,...], epsilon=1e-6)
        self.assertMatrixEqual(P[1,...], P1[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,
            observation_model,
            observation_noise,
            measurement)

        self.assertSequenceEqual(m.shape, (3,2,1))
        self.assertSequenceEqual(P.shape, (3,2,2))

        self.assertMatrixEqual(m[1,...], m1[1,...])
        self.assertMatrixEqual(m[2,...], mean[2,...], epsilon=1e-6)
        self.assertMatrixEqual(P[1,...], P1[1,...])

def test_update_with_nan_check(self):
        prior_mean = np.array([[1],[2],[3]])
        prior_covariance = np.eye(3)*2

        observation_model = np.ones((2,3))
        observation_noise = np.eye(2)*0.1

        measurement = np.array([[3],[np.nan]])

        m, P = primitives.update_with_nan_check(
            prior_mean,
            prior_covariance,
            observation_model,
            observation_noise,
            measurement)

        self.assertSequenceEqual(m.shape, (3,1))
        self.assertSequenceEqual(P.shape, (3,3))
        self.assertMatrixEqual(m, prior_mean)
        self.assertMatrixEqual(P, prior_covariance)

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,
            observation_model,
            observation_noise,
            measurement)

        self.assertSequenceEqual(m.shape, (1,1))
        self.assertSequenceEqual(P.shape, (1,1))