def test_message_to_parent(self):

        def check(indices, plates, shape, axis=-1, use_mask=False):
            mu = np.random.rand(*(plates+shape))
            alpha = np.random.rand(*(plates+shape))
            X = GaussianARD(mu, alpha, shape=shape, plates=plates)
            Y = Take(X, indices, plate_axis=axis)
            Z = GaussianARD(Y, 1, shape=shape)
            z = np.random.randn(*(Z.get_shape(0)))
            if use_mask:
                mask = np.mod(np.reshape(np.arange(np.prod(Z.plates)), Z.plates), 2) != 0
            else:
                mask = True
            Z.observe(z, mask=mask)
            X.update()
            (x0, x1) = X.get_moments()

            # For comparison, build the same model brute force
            X = GaussianARD(mu, alpha, shape=shape, plates=plates)

            # Number of trailing plate axes before the take axis
            N = len(X.plates) + axis

            # Reshape the take axes into a single axis
            z_shape = X.plates[:axis] + (-1,)
            if axis < -1:
                z_shape = z_shape + X.plates[(axis+1):]
            z_shape = z_shape + shape
            z = np.reshape(z, z_shape)

            # Reshape the take axes into a single axis
            if use_mask:
                mask_shape = X.plates[:axis] + (-1,)
                if axis < -1:
                    mask_shape = mask_shape + X.plates[(axis+1):]
                mask = np.reshape(mask, mask_shape)

            for (j, i) in enumerate(range(np.size(indices))):
                ind = np.array(indices).flatten()[i]
                index_x = N*(slice(None),) + (ind,)
                index_z = N*(slice(None),) + (j,)
                # print(index)
                Xi = X[index_x]
                zi = z[index_z]
                Zi = GaussianARD(Xi, 1, ndim=len(shape))
                if use_mask:
                    maski = mask[index_z]
                else:
                    maski = True
                Zi.observe(zi, mask=maski)

            X.update()

            self.assertAllClose(
                x0,
                X.get_moments()[0],
            )

            self.assertAllClose(
                x1,
                X.get_moments()[1],
            )

            return


        # Test scalar index
        check(1, (2,), ())
        check(1, (2, 3), ())
        check(1, (2, 3), (4,))
        check(1, (2, 3), (), axis=-2)
        check(1, (2, 3), (4,), axis=-2)
        check(1, (2,), (), use_mask=True)
        check(1, (2, 3), (), use_mask=True)
        check(1, (2, 3), (4,), use_mask=True)
        check(1, (2, 3), (), axis=-2, use_mask=True)
        check(1, (2, 3), (4,), axis=-2, use_mask=True)

        # Test vector index
        check([1, 1, 0, 1], (2,), ())
        check([1, 1, 0, 1], (2, 3), ())
        check([1, 1, 0, 1], (2, 3), (4,))
        check([1, 1, 0, 1], (2, 3), (), axis=-2)
        check([1, 1, 0, 1], (2, 3), (4,), axis=-2)
        check([1, 1, 0, 1], (2,), (), use_mask=True)
        check([1, 1, 0, 1], (2, 3), (), use_mask=True)
        check([1, 1, 0, 1], (2, 3), (4,), use_mask=True)
        check([1, 1, 0, 1], (2, 3), (), axis=-2, use_mask=True)
        check([1, 1, 0, 1], (2, 3), (4,), axis=-2, use_mask=True)

        # Test matrix index
        check([[1, 1, 2], [0, 2, 1]], (4,), ())
        check([[1, 1, 2], [0, 2, 1]], (4, 5), ())
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (6,))
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (), axis=-2)
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (6,), axis=-2)
        check([[1, 1, 2], [0, 2, 1]], (4,), (), use_mask=True)
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (), use_mask=True)
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (6,), use_mask=True)
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (), axis=-2, use_mask=True)
        check([[1, 1, 2], [0, 2, 1]], (4, 5), (6,), axis=-2, use_mask=True)

        pass
Test parent message computation in Take node
dataset/ETHPy150Open bayespy/bayespy/bayespy/inference/vmp/nodes/tests/test_take.py/TestTake.test_message_to_parent