update mx.random.normal calls

[LeonEricsson]

Proposed changes

updated existing mx.random.normal to make use of #638

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

Looks great, thanks!

[awni]

Any reason to keep this a draft?

[LeonEricsson]

Any reason to keep this a draft?

was waiting to see if #638 changed. do you mind taking a look at the prompt lookup pr when you've got some time over, its ready for merge

[awni]

👍