def METHOD_NAME():
    # test that gp proj gives expected shapes
    batch = 1
    hidden_size = 3
    dim = 4
    rank = 2
    states = mx.ndarray.ones(shape=(batch, dim, hidden_size))
    gp_proj = GPArgProj(rank=rank)
    gp_proj.initialize()
    distr_args = gp_proj(states)
    mu, D, W = distr_args
    assert mu.shape == (batch, dim)
    assert D.shape == (batch, dim)
    assert W.shape == (batch, dim, rank)
    distr = LowrankMultivariateGaussian(dim, rank, *distr_args)
    assert distr.mean.shape == (batch, dim)