Replace python random with torch.rand to enable dynamo.export

[BowenBao]

What does this PR do?

Related and Fixes pytorch/pytorch#102794

TL;DR dynamo graph breaks on python random.uniform(0, 1). The graph break can be prevented by replacing with torch.randn([]).

Example repro script

import torch
import torch._dynamo

from transformers import AutoTokenizer, BartForCausalLM

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
model = BartForCausalLM.from_pretrained("facebook/bart-base", add_cross_attention=False)
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
outputs = model(**inputs)

torch._dynamo.export(model, return_dict=False, **inputs)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sgugger]

You are touching multiple Flax model which shouldn't depend on torch, could you revert that?

[BowenBao]

You are touching multiple Flax model which shouldn't depend on torch, could you revert that?

Nice catch! Done.

[sgugger]

Thanks! @amyeroberts could you have a quick second look here?

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[amyeroberts]

Thanks for fixing this!

Changes LGTM 👍 I think technically this is slightly different as random.uniform can return 1, but for all practical purposes it doesn't change things.

[anijain2305]

@BowenBao Does this really solve the export problem. We are seeing export issues here - pytorch/pytorch#107587

If one peeks at the tensor value for the conditional, its a legit dynamic control flow. We might have to use torch.where or torch.cond.

[BowenBao]

@anijain2305 it does for inference export, since the actual condition is short circuited by self.training being False. Without the change, random.uniform(0, 1) leads to a graph break, although the value is unused.