diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index 0b18d9b06..f7eee95f7 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -233,8 +233,15 @@ def create_linear_map(signed=True, total_bits=8, add_zero=True): l = values.numel()//2 return torch.Tensor(values[:l].tolist() + [0]*gap + values[l:].tolist()) + def create_normal_map(offset=0.9677083, use_extra_value=True): - from scipy.stats import norm + try: + from scipy.stats import norm + except ImportError as ie: + raise ImportError( + "Scipy is required for `create_normal_map`. " + "Install `bitsandbytes` with the `[test]` extra." + ) from ie if use_extra_value: # one more positive value, this is an asymmetric type diff --git a/setup.py b/setup.py index b55678bf4..7e66508d0 100644 --- a/setup.py +++ b/setup.py @@ -26,11 +26,13 @@ def read(fname): license="MIT", keywords="gpu optimizers optimization 8-bit quantization compression", url="https://github.com/TimDettmers/bitsandbytes", - install_requires=['scipy'], packages=find_packages(), package_data={"": libs}, - install_requires=['torch', 'numpy', 'scipy'], - extras_require={'benchmark': ['pandas', 'matplotlib']}, + install_requires=['torch', 'numpy'], + extras_require={ + 'benchmark': ['pandas', 'matplotlib'], + 'test': ['scipy'], + }, long_description=read("README.md"), long_description_content_type="text/markdown", classifiers=[