[MPS] Fix batchnorm forward and backward pass

[DenisVieriu97]

Fixes batchnorm forward/backward pass and layer_norm:

Batchnorm Forward pass:

- fix batch_norm_mps_out key
- return 1/sqrt(var+epsilon) instead of var
- return empty tensor for mean and var if train is not enabled
- remove native_batch_norm from block list

Batchnorm Backward pass:

- add revert caculation for save_var used in backward path
- add backward test for native_batch_norm and _native_batch_norm_legit

Layer norm:

- remove the duplicate calculation from layer_norm_mps
- enable native_layer_norm backward test
- raise atol rtol for native_layer_norm

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/94351

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm jobs fail to access AMD apt repo

✅ No Failures

As of commit 604b8ad:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[razarmehr]

LGTM

[DenisVieriu97]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

[DenisVieriu97]

@pytorchbot merge -g

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[malfet]

Is there a reason not to use reciprocalWithTensor: here instead? It should be faster, shouldn't it? Or does it come with performance implications?
(PR is coming)

[kulinseth]

With the precise mode being default, there won't be any major perf implications.