AutoNormal and AutoLowRankMultivariateNormal can't process high dimensional dirichlet distribution?

There is an RV named phi with dim=5000 in my model. When I tried to build a guide using AutoNormal or AutoLowRankMultivariateNormal. I got the following error:

Error while computing log_prob at site 'phi_songId':
Expected value argument (Tensor of shape (6, 5000)) to be within the support (Simplex()) of the distribution Dirichlet(concentration: torch.Size([6, 5000])), but found invalid values:
tensor([[0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002],
        [0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002],
        [0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002],
        [0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002],
        [0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002],
        [0.0002, 0.0002, 0.0002,  ..., 0.0002, 0.0002, 0.0002]],
       device='cuda:1', grad_fn=<ExpandBackward0>)
  Trace Shapes:         
   Param Sites:         
  Sample Sites:         
phi_songId dist 6 | 5000
          value 6 | 5000

What can I do to solve this problem?

if you want useful recommendations you need to provide a lot more details. what prior distribution is phi governed by? etc

The prior distribution of phi_songId is Dirichlet with a symmetric prior parameter eta=torch.ones((5000, )). But I found that the symmetric prior parameter will lead to an unstable optimization result, so I randomly specified some dims of eta to tensor(50.) rather than tensor(1.).

it may help to use double precision. you may also be using too large a learning rate which is driving the parameters to unstable values

Thanks for your reply! I’ll try these solutions.

see here for additional things to try

I have already tried to use double precision. It worked but cost more GPU memory (I only have a 4*2080Ti server to train my model :smiling_face_with_tear:). Maybe I have to consider buying a new server or an account for Google colab pro+. Overall, thanks for your help. I really appreciate it. :kissing_heart: