import importlib


def extension_context(extension_name='cpu', **kw):
    """Get the context of the specified extension.

    All extension's module must provide `context(**kw)` function.

    Args:
        extension_name (str) : Module path relative to `nnabla.extensions`.
        kw (dict) : Additional keyword arguments for context function in a extension module.

    Returns:
        :class:`nnabla.Context`: The current extension context.

    Example:

        .. code-block:: python

            ctx = extension_context('cuda.cudnn', device_id=0)
            nn.set_default_context(ctx)

    """
    try:
        mod = importlib.import_module(
            '.' + extension_name, 'nnabla.extensions')
    except ImportError:
        mod = importlib.import_module('.' + extension_name, 'nnabla_ext')
    return mod.context(**kw)