`TestRsampleDeviceXPU.test_dirichlet_on_diagonal_xpu` fails on XPU

### 🐛 Describe the bug

The PR https://github.com/pytorch/pytorch/pull/195730 enables `distributions/test_distributions.py` tests on XPU. One of the tests (previously only executed on CPU), i.e., `TestRsampleDeviceXPU.test_dirichlet_on_diagonal` fails with the following error:

```
Traceback (most recent call last):
File "/usr/lib/python3.12/unittest/case.py", line 58, in testPartExecutor
yield
File "/usr/lib/python3.12/unittest/case.py", line 634, in run
self._callTestMethod(testMethod)
File "/usr/lib/python3.12/unittest/case.py", line 589, in _callTestMethod
if method() is not None:
^^^^^^^^
File "/home/gta/pytorch/torch/testing/_internal/common_utils.py", line 3886, in wrapper
method(*args, **kwargs)
File "/home/gta/pytorch/torch/testing/_internal/common_device_type.py", line 672, in instantiated_test
result = test(self, **param_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/gta/pytorch/test/distributions/test_distributions.py", line 4813, in test_dirichlet_on_diagonal
self.assertLess(
File "/usr/lib/python3.12/unittest/case.py", line 1257, in assertLess
self.fail(self._formatMessage(msg, standardMsg))
File "/usr/lib/python3.12/unittest/case.py", line 715, in fail
raise self.failureException(msg)
AssertionError: np.float64(0.009022388618624757) not less than 0.001 : Bad gradient dx[0]/dalpha[0] for Dirichlet([10.0, 10.0, 0.1])
x [0.31098405 0.4035269 0.40723997 0.48366055 0.48393574 0.4877867
0.49616152 0.50879043 0.54577863 0.54739094 0.5705272 0.5769548
0.59480554 0.60142463 0.6029198 0.6041471 0.6361284 0.6484746
0.6801623 0.75043166]
expected [0.02770732 0.02755608 0.02751564 0.02618291 0.0261765 0.02608564
0.02588079 0.02555343 0.02447226 0.02442113 0.02365215 0.02342704
0.02277651 0.022526 0.02246872 0.02242153 0.0211332 0.02060658
0.0191835 0.01568307]
actual [0.02770811 0.02755064 0.02750991 0.02633455 0.02594032 0.02608301
0.02587163 0.02544779 0.02446537 0.02441418 0.02364492 0.02342012
0.02277 0.02251964 0.0224624 0.02241528 0.02112807 0.02060197
0.01918039 0.01568398]
rel error [2.84653436e-05 1.97472654e-04 2.08579400e-04 5.79173560e-03
9.02238862e-03 1.00781202e-04 3.53998149e-04 4.13416022e-03
2.81372252e-04 2.84407664e-04 3.05390321e-04 2.95330383e-04
2.85963090e-04 2.82080142e-04 2.81605947e-04 2.78521520e-04
2.42859744e-04 2.23600025e-04 1.62287648e-04 5.83545417e-05]
max error 0.009022388618624757
at x=0.48393574357032776
```

### Versions

Pytorch: changes from https://github.com/pytorch/pytorch/pull/195730
torch-xpu-ops: main

Start in test/distributions/test_distributions.py at TestRsampleDeviceXPU.test_dirichlet_on_diagonal and reproduce the failure on XPU, comparing it with the CPU result and the tolerance assertion. Read the changes from PyTorch PR 195730 and the relevant torch-xpu-ops implementation to identify the source of the gradient discrepancy. Done means the targeted XPU test passes without masking unrelated failures.

Written by the indexing model from the issue text.