[release/2.14] Autograd Jacobian mismatch (GradcheckError) in RelaxedOneHotCategorical on XPU

Cases:
op_ut,third_party.torch-xpu-ops.test.xpu.test_distributions_xpu.TestDistributionsXPU,test_relaxed_one_hot_categorical_2d_xpu

## Root Cause

The test fails with `torch.autograd.gradcheck.GradcheckError` — a Jacobian mismatch between the numerical and analytical gradients of `RelaxedOneHotCategorical.log_prob()`. The numerical Jacobian produces `nan` while the analytical Jacobian returns a finite value (`-57982700.1`), indicating that the numerical differentiation step encounters a numerical instability (likely division by a value very close to zero or log of a near-zero value) that doesn't affect the analytical gradient path.

This is distinct from the other distribution failures: it's not a statistical sampling issue or a constraint violation, but an autograd correctness problem specific to the Gumbel-Softmax / RelaxedOneHotCategorical reparameterization on XPU.

## Reproducer

```bash
source ~/.bashrc && cd third_party/torch-xpu-ops/test/xpu && python -m pytest -sxv test_distributions_xpu.py -k "test_relaxed_one_hot_categorical_2d_xpu"
```

## Representative Error

```
=================================== FAILURES ===================================
_________ TestDistributionsXPU.test_relaxed_one_hot_categorical_2d_xpu _________
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/Projects/pytorch/torch/testing/_internal/common_utils.py", line 3591, in wrapper
method(*args, **kwargs)
File "/home/gta/Projects/pytorch/torch/testing/_internal/common_device_type.py", line 571, in instantiated_test
raise rte
File "/home/gta/Projects/pytorch/torch/testing/_internal/common_device_type.py", line 551, in instantiated_test
result = test(self, **param_kwargs)
File "/home/gta/Projects/pytorch/torch/testing/_internal/common_device_type.py", line 1637, in efail_fn
return fn(slf, *args, **kwargs)
File "/home/gta/Projects/pytorch/torch/testing/_internal/common_utils.py", line 2753, in wrapper
return fn(test_self, *args, **kwargs)
File "test/distributions/test_distributions.py", line 2232, in test_relaxed_one_hot_categorical_2d
self._gradcheck_log_prob(
File "test/distributions/test_distributions.py", line 1364, in _gradcheck_log_prob
gradcheck(apply_fn, (s,) + tuple(ctor_params), raise_exception=True)
File "/home/gta/Projects/pytorch/torch/autograd/gradcheck.py", line 2104, in gradcheck
return _gradcheck_helper(**args)
File "/home/gta/Projects/pytorch/torch/autograd/gradcheck.py", line 2133, in _gradcheck_helper
_gradcheck_real_imag(
File "/home/gta/Projects/pytorch/torch/autograd/gradcheck.py", line 1517, in _gradcheck_real_imag
gradcheck_fn(
File "/home/gta/Projects/pytorch/torch/autograd/gradcheck.py", line 1974, in _fast_gradcheck
_check_analytical_numerical_equal(
File "/home/gta/Projects/pytorch/torch/autograd/gradcheck.py", line 1902, in _check_analytical_numerical_equal
raise GradcheckError(
torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for output 0 with respect to input 0,
numerical:nan
analytical:-57982700.11608397
```