feat: throwing a warning when a sub-backend implementation is available but is not being used.

ivy/__init__.py

@@ -1199,6 +1199,7 @@ def dynamic_backend_as(value):
     unset_sub_backend,
     clear_sub_backends,
     available_sub_backends,
+    available_sub_backend_implementations,
 )
 
 

ivy/func_wrapper.py

@@ -1580,6 +1580,38 @@ def func(x):
     return _handle_backend_invalid
 
 
+def _handle_efficient_implementation_available(fn: Callable) -> Callable:
+    @functools.wraps(fn)
+    def _wrapper(*args, **kwargs):
+        """
+        Throws a warning whenever the function is called, indicating that an efficient
+        implementation is available. This should be used for functions which have an
+        efficient implementation in a sub backend.
+
+        Parameters
+        ----------
+        args
+            The arguments to be passed to the function.
+
+        kwargs
+            The keyword arguments to be passed to the function.
+
+        Returns
+        -------
+            The return of the function.
+        """
+        ivy.warn(
+            f"An efficient implementation of {fn.__name__} is available "
+            "in these sub backends: "
+            f"{ivy.available_sub_backend_implementations(fn.__name__)}.\n"
+            "use ivy.set_sub_backend('<sub_backend_name>') to use it."
+        )
+        return fn(*args, **kwargs)
+
+    _wrapper.efficient_implementation_available = True
+    return _wrapper
+
+
 attribute_dict = {
     "unsupported_dtypes",
     "supported_dtypes",

ivy/utils/backend/handler.py

@@ -10,7 +10,7 @@
 from ivy.utils import _importlib, verbosity
 
 # local
-from ivy.func_wrapper import _wrap_function
+from ivy.func_wrapper import _wrap_function, _handle_efficient_implementation_available
 from ivy.utils.backend.sub_backend_handler import _clear_current_sub_backends
 from ivy.utils.exceptions import _handle_inplace_mode
 
@@ -243,16 +243,17 @@ def current_backend(*args, **kwargs):
     return importlib.import_module(_backend_dict[implicit_backend])
 
 
-def _set_backend_as_ivy(
+def _modify_module_symbols(
     original_dict, target, backend, invalid_dtypes=None, backend_str=None
 ):
     invalid_dtypes = (
         backend.invalid_dtypes if invalid_dtypes is None else invalid_dtypes
     )
     backend_str = backend.current_backend_str() if backend_str is None else backend_str
     for k, v in original_dict.items():
+        v = _wrap_if_got_efficient_implementations(v) if v else v
         compositional = k not in backend.__dict__
-        if k not in backend.__dict__:
+        if compositional:
             if k in invalid_dtypes and k in target.__dict__:
                 del target.__dict__[k]
                 continue
@@ -265,7 +266,7 @@ def _set_backend_as_ivy(
             and "ivy.functional." in v.__name__
             and os.path.join("{}", "__init__.py").format(backend_str) not in v.__file__
         ):
-            _set_backend_as_ivy(
+            _modify_module_symbols(
                 v.__dict__,
                 target.__dict__[k],
                 backend.__dict__[k],
@@ -274,6 +275,13 @@ def _set_backend_as_ivy(
             )
 
 
+def _wrap_if_got_efficient_implementations(v):
+    if callable(v):
+        if ivy.available_sub_backend_implementations(v.__name__):
+            return _handle_efficient_implementation_available(v)
+    return v
+
+
 def _handle_backend_specific_vars(target, backend):
     if backend.current_backend_str() == "numpy":
         target.set_default_device("cpu")
@@ -453,7 +461,7 @@ def set_backend(backend: str, dynamic: bool = False):
             ivy.set_global_attr("RNG", ivy.functional.backends.jax.random.RNG)
         backend_stack.append(backend)
         set_backend_to_specific_version(backend)
-        _set_backend_as_ivy(ivy_original_dict, ivy, backend)
+        _modify_module_symbols(ivy_original_dict, ivy, backend)
         # following snippet is required to update the ivy.functional namespace with
         # backend-specific functions
         for key, _ in ivy.__dict__.items():
@@ -640,7 +648,7 @@ def with_backend(backend: str, cached: bool = True):
         set_backend_to_specific_version(backend_module)
         # We know for sure that the backend stack is empty
         # no need to do backend unsetting
-        ivy_pack.utils.backend.handler._set_backend_as_ivy(
+        ivy_pack.utils.backend.handler._modify_module_symbols(
             ivy_pack.__dict__.copy(), ivy_pack, backend_module
         )
         # TODO use a refactored code from ivy.set_backend

ivy/utils/backend/sub_backend_handler.py

@@ -9,53 +9,62 @@
 
 
 _backends_subpackage_path = "ivy.functional.backends"
-_sub_backend_dict = dict()
-_backend_to_sub_backends_dict = dict()
-
-# dynamic sub_backend detection
-for backend in os.listdir(
-    os.path.join(
-        ivy.__path__[0].rpartition(os.path.sep)[0],  # type: ignore
-        _backends_subpackage_path.replace(".", os.path.sep),
-    )
-):
-    if not backend[0].isalpha():
-        continue
-
-    sub_backends_dir = os.path.join(
-        ivy.__path__[0].rpartition(os.path.sep)[0],
-        _backends_subpackage_path.replace(".", os.path.sep),
-        backend,
-        "sub_backends",
-    )
-    for sub_backend in os.listdir(sub_backends_dir):
-        if not sub_backend[0].isalpha():
+_sub_backend_dict: dict[str, str] = dict()
+_backend_to_sub_backends_dict: dict[str, list] = dict()
+_available_sub_backends_implementations_dict: dict[str, dict[str, list]] = dict()
+
+
+def _detect_sub_backends_dynamically():
+    for backend in os.listdir(
+        os.path.join(
+            ivy.__path__[0].rpartition(os.path.sep)[0],  # type: ignore
+            _backends_subpackage_path.replace(".", os.path.sep),
+        )
+    ):
+        if not backend[0].isalpha():
             continue
-        _sub_backend_dict[sub_backend] = (
-            f"{_backends_subpackage_path}.{backend}.sub_backends.{sub_backend}"
+
+        sub_backends_dir = os.path.join(
+            ivy.__path__[0].rpartition(os.path.sep)[0],
+            _backends_subpackage_path.replace(".", os.path.sep),
+            backend,
+            "sub_backends",
         )
-        try:
-            _backend_to_sub_backends_dict[backend].append(sub_backend)
-        except KeyError:
-            _backend_to_sub_backends_dict[backend] = [sub_backend]
+        for sub_backend in os.listdir(sub_backends_dir):
+            if not sub_backend[0].isalpha():
+                continue
+            _sub_backend_dict[sub_backend] = (
+                f"{_backends_subpackage_path}.{backend}.sub_backends.{sub_backend}"
+            )
+            try:
+                _backend_to_sub_backends_dict[backend].append(sub_backend)
+            except KeyError:
+                _backend_to_sub_backends_dict[backend] = [sub_backend]
+
+
+_detect_sub_backends_dynamically()
+
 
+def _get_all_sub_backends():
+    sub_backends = []
+    for v in _backend_to_sub_backends_dict.values():
+        sub_backends.extend(v)
+    return sub_backends
 
-_all_sub_backends = []
 
-for v in _backend_to_sub_backends_dict.values():
-    _all_sub_backends.extend(v)
+_all_sub_backends = _get_all_sub_backends()
 
 
 original_backend_dict = None
 
 
 def set_sub_backend(sub_backend_str: str):
     if ivy.backend == "":
-        logging.warn("You must set a backend first")
+        logging.warning("You must set a backend first")
         return
 
     if ivy.current_backend_str() not in _backend_to_sub_backends_dict.keys():
-        logging.warn(
+        logging.warning(
             f"backend {ivy.current_backend_str()} does not have any"
             " supported sub_backends"
         )
@@ -68,7 +77,7 @@ def set_sub_backend(sub_backend_str: str):
         )
 
     if sub_backend_str not in _backend_to_sub_backends_dict[ivy.current_backend_str()]:
-        logging.warn(
+        logging.warning(
             f"{ivy.current_backend_str()} does not support"
             f" {sub_backend_str} as a sub_backend"
         )
@@ -87,7 +96,7 @@ def set_sub_backend(sub_backend_str: str):
     ivy.current_backend().sub_backends._current_sub_backends.append(sub_backend_str)
 
 
-# this is very similiar to _set_backend_as_ivy in handler.py, with a minor change
+# this is very similar to _set_backend_as_ivy in handler.py, with a minor change
 def _set_sub_backend_as_ivy(
     original: dict, target: ModuleType, sub_backend: ModuleType
 ):
@@ -200,3 +209,67 @@ def find_available_sub_backends(sub_backends_loc):
             available_sub_backends.append(sub_backend)
 
     return available_sub_backends
+
+
+def _find_available_sub_backend_implementations(sub_backends):
+    result = dict()
+    for sub in sub_backends:
+        try:
+            sub_backend = ivy.utils.dynamic_import.import_module(_sub_backend_dict[sub])
+        except ModuleNotFoundError:
+            continue
+        for k, v in sub_backend.__dict__.items():
+            if callable(v) and not k.startswith("_"):
+                result[k] = result.get(k, []) + [sub]
+
+    return result
+
+
+def available_sub_backend_implementations(fn_name: str) -> list:
+    """
+    Return whether a sub-backend implementation is available for `obj`.
+
+    Parameters
+    ----------
+    fn_name : str
+        the object for which to check if a sub-backend implementation is available.
+
+    Returns
+    -------
+    ret : list
+        a list of sub-backend implementations available for `obj`.
+
+    Examples
+    --------
+    >>> import ivy
+    >>> ivy.set_backend('torch')
+    >>> ivy.available_sub_backend_implementations("scaled_dot_product_attention")
+    ['xformers']
+    >>> ivy.set_backend('numpy')
+    >>> ivy.available_sub_backend_implementations("scaled_dot_product_attention")
+    []
+    """
+    sub_backends = ivy.current_backend().available_sub_backends()
+    if not sub_backends:
+        return []
+    if not _sub_backends_implementations_already_verified():
+        _verify_sub_backends_implementations(sub_backends)
+    return _available_implementations_for(fn_name)
+
+
+def _sub_backends_implementations_already_verified():
+    return (
+        ivy.current_backend_str() in _available_sub_backends_implementations_dict.keys()
+    )
+
+
+def _verify_sub_backends_implementations(sub_backends):
+    _available_sub_backends_implementations_dict[ivy.current_backend_str()] = (
+        _find_available_sub_backend_implementations(sub_backends)
+    )
+
+
+def _available_implementations_for(fn_name):
+    return _available_sub_backends_implementations_dict[ivy.current_backend_str()].get(
+        fn_name, []
+    )

ivy_tests/test_ivy/test_misc/test_func_wrapper.py

@@ -209,3 +209,9 @@ def test_views(array_to_update, backend_fw):
     assert np.allclose(d, d_copy + 1)
     assert np.allclose(e[0], e_copy + 1)
     ivy.previous_backend()
+
+
+def test_warn_efficient_implementations():
+    to_test = ivy.func_wrapper._handle_efficient_implementation_available(ivy.array)
+    with pytest.warns(UserWarning):
+        to_test(1)

ivy_tests/test_ivy/test_misc/test_sub_backends_available.py

@@ -0,0 +1,65 @@
+import pytest
+
+import ivy
+from ivy.utils.backend.sub_backend_handler import (
+    _find_available_sub_backend_implementations,
+)
+
+
+@pytest.fixture
+def setup(monkeypatch):
+    def _fn(_):
+        return {"scaled_dot_product_attention": ["xformers"]}
+
+    monkeypatch.setattr(
+        "ivy.utils.backend.sub_backend_handler."
+        "_find_available_sub_backend_implementations",
+        _fn,
+    )
+
+
+def test_find_available_sub_backend_implementations():
+    ivy.set_backend("torch")
+    result = _find_available_sub_backend_implementations(["xformers"])
+    assert isinstance(result, dict)
+    assert all(isinstance(key, str) for key in result.keys())
+    assert all(isinstance(value, list) for value in result.values())
+
+
+def test_no_warning_when_no_sub_backend_implementation_available():
+    ivy.set_backend("numpy")
+    q = ivy.array([[[0.2, 1.0], [2.2, 3.0], [4.4, 5.6]]])
+    k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]])
+    v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]])
+    with pytest.warns(None) as record:
+        ivy.scaled_dot_product_attention(
+            q, k, v, scale=1, dropout_p=0.1, is_causal=True, training=True
+        )
+    assert len(record) == 0
+
+
+def test_sub_backend_implementation_available(setup):
+    ivy.set_backend("torch")
+    sub_backends = ivy.available_sub_backend_implementations(
+        "scaled_dot_product_attention"
+    )
+    assert "xformers" in sub_backends
+
+
+def test_sub_backend_implementation_not_available():
+    ivy.set_backend("numpy")
+    sub_backends = ivy.available_sub_backend_implementations(
+        "scaled_dot_product_attention"
+    )
+    assert not sub_backends
+
+
+def test_throw_warning_when_sub_backend_implementation_available_but_not_used(setup):
+    ivy.set_backend("torch")
+    q = ivy.array([[[0.2, 1.0], [2.2, 3.0], [4.4, 5.6]]])
+    k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]])
+    v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]])
+    with pytest.warns(UserWarning):
+        ivy.scaled_dot_product_attention(
+            q, k, v, scale=1, dropout_p=0.1, is_causal=True, training=True
+        )