Repository: pytorch
Issue: [bug] Internal assert failed when using pyro
Labels: module: distributions,triaged,module: macos,module: linear algebra,module: python frontend

### 🐛 Describe the bug

The initialization of a wishart distribution below fails. Reproducing example:

```python
import pyro, torch
torch.manual_seed(42)

print(pyro.__version__)  # 1.8.4
print(torch.__version__) # 2.0.0

n = 5000

X = pyro.param(f'X', torch.randn(n, 2))
log_s = pyro.param(f's', torch.randn(1))
M_mean = X @ X.T + log_s.exp() * torch.eye(n)
pyro.distributions.Wishart(n, M_mean)
```

The error that is thrown:
```python
** On entry to SSYEVD, parameter number  8 had an illegal value
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In [1], line 12
     10 log_s = pyro.param(f's', torch.randn(1))
     11 M_mean = X @ X.T + log_s.exp() * torch.eye(n)
---> 12 pyro.distributions.Wishart(n, M_mean)

File ~/miniconda3/lib/python3.9/site-packages/pyro/distributions/distribution.py:24, in DistributionMeta.__call__(cls, *args, **kwargs)
     22     if result is not None:
     23         return result
---> 24 return super().__call__(*args, **kwargs)

File ~/miniconda3/lib/python3.9/site-packages/torch/distributions/wishart.py:109, in Wishart.__init__(self, df, covariance_matrix, precision_matrix, scale_tril, validate_args)
    106 if self.df.lt(event_shape[-1]).any():
    107     warnings.warn("Low df values detected. Singular samples are highly likely to occur for ndim - 1 < df < ndim.")
--> 109 super().__init__(batch_shape, event_shape, validate_args=validate_args)
    110 self._batch_dims = [-(x + 1) for x in range(len(self._batch_shape))]
    112 if scale_tril is not None:

File ~/miniconda3/lib/python3.9/site-packages/torch/distributions/distribution.py:60, in Distribution.__init__(self, batch_shape, event_shape, validate_args)
     58     continue  # skip checking lazily-constructed args
     59 value = getattr(self, param)
---> 60 valid = constraint.check(value)
     61 if not valid.all():
     62     raise ValueError(
     63         f"Expected parameter {param} "
     64         f"({type(value).__name__} of shape {tuple(value.shape)}) "
   (...)
     67         f"but found invalid values:\n{value}"
     68     )

File ~/miniconda3/lib/python3.9/site-packages/pyro/distributions/torch_patch.py:81, in _PositiveDefinite_check(self, value)
     78 batch_shape = value.shape[:-2]
     79 flattened_value = value.reshape((-1,) + matrix_shape)
     80 return torch.stack(
---> 81     [torch.linalg.eigvalsh(v)[:1] > 0.0 for v in flattened_value]
     82 ).view(batch_shape)

File ~/miniconda3/lib/python3.9/site-packages/pyro/distributions/torch_patch.py:81, in <listcomp>(.0)
     78 batch_shape = value.shape[:-2]
     79 flattened_value = value.reshape((-1,) + matrix_shape)
     80 return torch.stack(
---> 81     [torch.linalg.eigvalsh(v)[:1] > 0.0 for v in flattened_value]
     82 ).view(batch_shape)

RuntimeError: false INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/BatchLinearAlgebra.cpp":1539, please report a bug to PyTorch. linalg.eigh: Argument 8 has illegal value. Most certainly there is a bug in the implementation calling the backend library.
```

### Versions

```
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.102)
CMake version: version 3.26.0
Libc version: N/A

Python version: 3.9.15 | packaged by conda-forge | (main, Nov 22 2022, 08:52:10)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-12.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2

Versions of relevant libraries:
[pip3] botorch==0.7.2
[pip3] gpytorch==1.9.0
[pip3] mypy==1.1.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] numpyro==0.10.1
[pip3] pytorch-lightning==1.7.7
[pip3] torch==2.0.0
[pip3] torch-cluster==1.6.0
[pip3] torch-geometric==2.2.0
[pip3] torch-scatter==2.1.0
[pip3] torch-sparse==0.6.16
[pip3] torch-spline-conv==1.2.1
[pip3] torchmetrics==0.10.2
[pip3] torchpwl==0.2.0
[pip3] torchsde==0.2.5
[pip3] torchvision==0.14.0
[conda] botorch                   0.7.2                    pypi_0    pypi
[conda] gpytorch                  1.9.0                    pypi_0    pypi
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] numpy-base                1.21.5           py39hadd41eb_3  
[conda] numpyro                   0.10.1                   pypi_0    pypi
[conda] pytorch-lightning         1.7.7                    pypi_0    pypi
[conda] torch                     2.0.0                    pypi_0    pypi
[conda] torch-cluster             1.6.0                    pypi_0    pypi
[conda] torch-geometric           2.2.0                    pypi_0    pypi
[conda] torch-scatter             2.1.0                    pypi_0    pypi
[conda] torch-sparse              0.6.16                   pypi_0    pypi
[conda] torch-spline-conv         1.2.1                    pypi_0    pypi
[conda] torchmetrics              0.10.2                   pypi_0    pypi
[conda] torchpwl                  0.2.0                    pypi_0    pypi
[conda] torchsde                  0.2.5                    pypi_0    pypi
[conda] torchvision               0.14.0                   pypi_0    pypi
```

cc @fritzo @neerajprad @alicanb @nikitaved @malfet @albanD @jianyuh @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @Lezcano