Repository: pytorch
Issue: Failed to Open libnvrtc-builtins.so.11.7
Labels: oncall: jit,module: cuda

### 🐛 Describe the bug

When running a model with a custom layer jitted with torchscript, the following error is thrown when attempting to use it via Rust bindings (e.g. tch):

```
panicked at 'called `Result::unwrap()` on an `Err` value: Torch("The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: false INTERNAL ASSERT FAILED at \"../torch/csrc/jit/codegen/cuda/executor_utils.cpp

--A BUNCH OF C++ TRASH--

CUDA NVRTC compile error: nvrtc: error: failed to open libnvrtc-builtins.so.11.7.  Make sure that libnvrtc-builtins.so.11.7 is installed correctly.")'
```

I've installed libtorch via the Pytorch download and it actually has the shared library `libnvrtc-builtins.so.11.7` but it's named `libnvrtc-builtins-7237cb5d.so.11.7` or some other thing. Anyways, the simple workaround is to literally rename the shared library to  `libnvrtc-builtins.so.11.7`, but I suppose this isn't desirable since there's  a reason it's named this.

The example model is huge but I was able to narrow it down to the layer that triggers this. I don't know if the layer, itself, triggers this, but if you bolt it onto a bigger model, it seems to trigger the tomfoolery described above:
```
from torch import Module
import math
import torch

class NormalOutputLayer(Module):
    """
    An output layer that takes in an 2N inputs which represent the means and standard deviations of n normal
    distributions, and outputs N elements that are sampled from those distributions.
    :param samples: Number of samples to draw.
    :param reduce: Reduction method. Can be "mean," "mean_std,", "median," "49th", "49th_std", "sum," or None.
                   Default is None.
    """

    def __init__(self, samples: int = 1, reduce: str = None):
        super(NormalOutputLayer, self).__init__()

        if samples <= 0:
            raise ValueError(f"Samples must be an integer greater than 0")

        if (
            reduce is not None
            and reduce != "mean"
            and reduce != "sum"
            and reduce != "mean_std"
            and reduce != "median"
            and reduce != "49th"
            and reduce != "49th_std"
        ):
            raise ValueError(f"Unknown reduction method: {reduce}")

        self.samples = samples
        self.reduce = reduce

    def forward(self, mu: torch.Tensor, sigma: torch.Tensor):
        """
        Dynamically generates normal distributions and samples from them. This layer will take the elements along the
        last axis and use the first half of the values as the mu values and the last half as the sigma values.
        :param mu: Tensor that contains the means.
        :param sigma: Tensor that contains the standard deviations.
        :return: N samples from a normal distribution that have been reduced appropriately.
        """

        if mu.size() != sigma.size():
            raise ValueError(
                f"mu tensor and sigma tensor are not equal: {mu.size()} vs {sigma.size()}"
            )

        # Squeeze batched input that has only one element.
        if len(mu.size()) == 2 and mu.size()[0] == 1:
            manually_batch_output = True
            mu = torch.squeeze(mu)
            sigma = torch.squeeze(sigma)
        else:
            manually_batch_output = False

        normal_dist = Normal(loc=mu, scale=sigma)
        # Re-parametrization trick (ref: https://stackoverflow.com/q/60533150/16070598)
        samples = normal_dist.rsample(sample_shape=[self.samples])

        if self.reduce == "mean":
            result = torch.mean(samples, dim=0)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce == "sum":
            result = torch.sum(samples, dim=0)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce == "mean_std":
            std, mean = torch.std_mean(samples, dim=0)
            result = torch.transpose(torch.stack((mean, std)), 0, 1)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce == "median":
            result = torch.quantile(samples, 0.5, dim=0)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce == "49th":
            result = torch.quantile(samples, 0.49, dim=0)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce == "49th_std":
            quantile = torch.quantile(samples, 0.49, dim=0)
            std = torch.std(samples, dim=0)
            result = torch.transpose(torch.stack((quantile, std)), 0, 1)
            return torch.unsqueeze(result, dim=0) if manually_batch_output else result
        elif self.reduce is None:
            return torch.unsqueeze(samples, dim=0) if manually_batch_output else samples
        else:
            raise ValueError(f"Unknown reduce method: {self.reduce}")
```

The `Normal` class is as follows:

```
import math
import torch

class Normal:
    """
    Creates a normal (also called Gaussian) distribution parameterized by
    :attr:`loc` and :attr:`scale`. This is a copy of PyTorch's
    torch.distributions.Normal class that is compatiable with PyTorch's
    JIT script operation.

    Example::

        >>> # xdoctest: +IGNORE_WANT("non-deterinistic")
        >>> m = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
        >>> m.sample()  # normally distributed with loc=0 and scale=1
        tensor([ 0.1046])

    Args:
        loc (float or Tensor): mean of the distribution (often referred to as mu)
        scale (float or Tensor): standard deviation of the distribution
            (often referred to as sigma)
    """

    @property
    def mean(self):
        return self.loc

    @property
    def stddev(self):
        return self.scale

    @property
    def variance(self):
        return self.stddev.pow(2)

    def __init__(self, loc: torch.Tensor, scale: torch.Tensor):
        resized_ = torch.broadcast_tensors(loc, scale)
        self.loc = resized_[0]
        self.scale = resized_[1]
        self._batch_shape = list(self.loc.size())

    def _extended_shape(self, sample_shape: list[int]) -> list[int]:
        return sample_shape + self._batch_shape

    def sample(self, sample_shape: list[int]) -> torch.Tensor:
        shape = self._extended_shape(sample_shape)
        return torch.normal(self.loc.expand(shape), self.scale.expand(shape))

    def rsample(self, sample_shape: list[int]) -> torch.Tensor:
        shape: list[int] = self._extended_shape(sample_shape)
        eps = torch.normal(
            torch.zeros(shape, device=self.loc.device),
            torch.ones(shape, device=self.scale.device),
        )
        return self.loc + eps * self.scale

    def log_prob(self, value: torch.Tensor) -> torch.Tensor:
        var = self.scale**2
        log_scale = self.scale.log()
        return (
            -((value - self.loc) ** 2) / (2 * var)
            - log_scale
            - math.log(math.sqrt(2 * math.pi))
        )

    def entropy(self) -> torch.Tensor:
        return 0.5 + 0.5 * math.log(2 * math.pi) + torch.log(self.scale)
```

### Versions

```
Collecting environment information...
PyTorch version: 1.13.0
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Pop!_OS 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.6 (main, Oct 24 2022, 16:07:47) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.0.12-76060006-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Super with Max-Q Design
Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.3
[pip3] torch==1.13.0
[pip3] torchaudio==0.13.0
[pip3] torchvision==0.14.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.3          py310hd5efca6_1  
[conda] numpy-base                1.23.3          py310h8e6c178_1  
[conda] pytorch                   1.13.0          py3.10_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.13.0              py310_cu116    pytorch
[conda] torchvision               0.14.0              py310_cu116    pytorch
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @ngimel