Merge branch 'review-borchify' into 'master'

        self.observed = Observed()

        self._used_rvs = set()

    @property

    def internal_modules(self):

        """Get the internal modules borch uses, like `prior`, `posterior`, `observed`"""

        return (self.posterior, self.prior, self.observed)

    def __deepcopy__(self, mmo):

        """Support `copy.deepcopy()`

"""

from borch.nn.torch_proxies import *

from borch.nn.compose import Sequential, Concurrent, Select

from borch.nn.reshape import Expand, ExpandBatch, View, BatchView, Flatten

from borch.nn.utility_modules import Identity, Apply, Cat

from borch.nn.recall import Recall

from borch.nn import utils

from borch.nn.borchify import borchify_module, borchify_network

from borch.module import Module

Functions to 'borchify' PyTorch modules/networks.

"""

from typing import Optional

from copy import deepcopy

from torch.nn import Module, modules

from torch.nn import Module

from borch import nn

from borch.module import Module as BorchModule

from borch.posterior import Posterior, Normal

from borch.rv_factories import apply_rv_factory

from borch.nn.torch_proxies import get_rv_factory

from borch.nn.torch_proxies import (

    get_rv_factory,

    _torch_proxy_class_factory,

    TORCH_BORCH_MAP,

    NOT_BORCHIFIED,

)

from borch.utils.module_utils import copy_module_attributes

def as_borch_module(cls, *args, existing=None, **kwargs):

def as_torch_module(cls, *args, existing=None, **kwargs):

    """Convert a module to a borch module"""

    obj = cls.__new__(cls, *args, **kwargs)

    if isinstance(existing, modules.Module):

    if existing is not None:

        copy_module_attributes(original=existing, new=obj)

        # Initialise only with Module. The corresponding ``torch.nn.XXX``

        # initialisations have already been performed (this is for

        # borchification purposes).

    return obj

def as_borch_module(cls, *args, existing=None, **kwargs):

    """Convert a module to a borch module"""

    obj = cls.__new__(cls, *args, **kwargs)

    BorchModule.__init__(obj, *args, **kwargs)

    if existing is not None:

        copy_module_attributes(original=existing, new=obj)

    else:

        # Initialise with super. This is a new object creation and all

        # super inits should be called.

        obj.__init__(*args, **kwargs)

    return obj

def _register_rvs_with_posterior(module):

    if hasattr(module, "prior"):

        for key in module.prior._modules.keys():  # pylint: disable=protected-access

            getattr(module, key, None)

_BORCHIFIED_CLASS_MAP = {}

def _get_borchified_class(cls):

    if cls in TORCH_BORCH_MAP:

        return TORCH_BORCH_MAP[cls]

    if cls not in _BORCHIFIED_CLASS_MAP:

        _BORCHIFIED_CLASS_MAP[cls] = _torch_proxy_class_factory(cls)

    return _BORCHIFIED_CLASS_MAP[cls]

def borchify_module(

    module: Module, rv_factory: Optional[callable] = None, posterior: Posterior = None

) -> BorchModule:

        >>> type(blinear)

        <class 'borch.nn.torch_proxies.Linear'>

    """

    # Unexpected keyword arg is consumed by the metaclass

    # pylint: disable=unexpected-keyword-arg

    cls_name = type(module).__name__

    if posterior is None:

        posterior = Normal()

    if isinstance(module, BorchModule):

        new = deepcopy(module)

        new.posterior = posterior

        _register_rvs_with_posterior(new)

        return new

    if cls_name not in dir(nn):

        # If a custom class which inherits from torch.nn.Module was created

        # then it does not exist in borch.nn (obviously). In this case we can

        # simply construct from Module. NB we also still apply the ``rv_factory``

        # to all parameters.

        new = as_borch_module(BorchModule, existing=module)

    else:

        new_module = getattr(nn, cls_name)

        if issubclass(new_module, BorchModule):

            new = as_borch_module(new_module, existing=module)

        else:  # pragma: no cover

            # TODO: it is not sufficient to only check the string to only check for

            # name above 'if cls_name not in dir(nn):' We should check if the actual

            # class is available in nn - see issue 295

            new = new_module()

            copy_module_attributes(module, new)

    cls_name = type(module).__name__

    new_module = _get_borchified_class(type(module))

    if new_module in NOT_BORCHIFIED:

        return as_torch_module(new_module, existing=module)

    new = as_borch_module(new_module, existing=module, posterior=posterior)

    temp_rv_factory = get_rv_factory(cls_name) if rv_factory is None else rv_factory

    new.posterior = posterior

    return apply_rv_factory(new, temp_rv_factory)

    apply_rv_factory(new, temp_rv_factory)

    _register_rvs_with_posterior(new)

    return new

def borchify_network(

    module: Module,

    rv_factory: Optional[callable] = None,

    posterior_creator: callable = Normal,

    posterior_creator: callable = None,

    cache: dict = None,

) -> BorchModule:

    """Borchify a whole network. This applies ``borchify_module`` recursively on

      >>> net = Net()

      >>> bnet = borchify_network(net)

      >>> type(bnet)

      <class 'borch.module.Module'>

      <class 'borch.nn.torch_proxies.Net'>

      >>> type(bnet.linear)

      <class 'borch.nn.torch_proxies.Linear'>

    """

    _id = id(module)

    if _id in cache:

        # it is implicit that all submodules already exist

        return cache[_id]

    if posterior_creator is None:

        posterior_creator = lambda: Normal(-4, loc_at_prior_mean=False)

    new = borchify_module(module, rv_factory, posterior_creator())

    cache[_id] = new

    for name, mod in module._modules.items():

        # NB we recurse into ``mod`` if it does not appear in ``cache``

        if isinstance(module, BorchModule) and mod in module.internal_modules:

            continue

        adding = cache.get(

            id(mod), borchify_network(mod, rv_factory, posterior_creator, cache)

        )

        new.add_module(name, adding)

    return new

"""

Modules for composition such as Sequential, Select and Concurrent. The modules have

no functionality per se, but they dictate how their children should be executed.

"""

from torch import nn

class Concurrent(nn.Module):

    """

    Module for splitting input x to different modules and return a tuple of

    the results

    Args:

        branches: The to feed the input through

    Examples:

        >>> import torch

        >>> split = Concurrent(nn.Identity(), nn.Identity())

        >>> x = torch.ones(2)

        >>> split(x)

        (tensor([1., 1.]), tensor([1., 1.]))

    """

    def __init__(self, *branches):

        super().__init__()

        self.branches = nn.ModuleList(branches)

    def forward(self, x):  # pylint: disable=arguments-differ

        """Forward method."""

        return tuple(branch(x) for branch in self.branches)

def apply_sequentially(x, modules):

    """

    Applies an iterable of modules sequentially to an input x.

    Example:

        >>> import torch

        >>> from borch import nn

        >>> test_img = torch.randn(1, 3, 20, 20)

        >>> m1 = nn.Conv2d(3, 5, 3)

        >>> m2 = nn.Conv2d(5, 9, 3)

        >>> m3 = nn.Conv2d(9, 3, 3)

        >>> modules = [m1, m2, m3]

        >>> out1 = apply_sequentially(test_img, modules)

        >>> out2 = m1(test_img)

        >>> out2 = m2(out2)

        >>> out2 = m3(out2)

        out1 and out2 has the same value whether we apply modules manually or use

        apply_sequentially

        >>> assert(torch.all(torch.eq(out1, out2)))

    """

    for mod in modules:

        if isinstance(x, (tuple, list)):

            x = mod(*x)

        else:

            x = mod(x)

    return x

class Sequential(nn.Sequential):

    # pylint: disable=arguments-differ, missing-docstring

    __doc__ = (

        nn.Sequential.__doc__

        + """

        Notes:

            **NOTE THAT ANYTHING WHICH INHERITS FROM SEQUENTIAL SHOULD NOT IMPLEMENT

            CUSTOM LOGIC IN THE FORWARD PASS.** The forward pass of a ``Sequential``

            will by design execute modules in the order they were added.

            Also, note that any input to a module which is a list or tuple will be

            expanded in the Sequential forward pass.

        """

    )

    def __getitem__(self, item):

        modules = list(self._modules.values())

        if isinstance(item, slice):

            return Sequential(*modules[item])

        return modules[item]

    def forward(self, *x):

        """Forward method."""

        return apply_sequentially(x, self.children())

class Select(nn.Module):

    """

    Selects a list of indices in an input tuple and passes them forward in a tuple,

    discarding the rest.

    Args:

        idxs: Indices to select from tuple. If the indices of the list itself is a list

          for example [1, [3, 0]] we will select indexable_structure[1] and

          indexable_structure[3][0].

    Examples

        >>> import torch

        >>> select = Select(0)

        >>> x = (torch.ones(2), torch.ones(2))

        >>> select(x)

        (tensor([1., 1.]),)

        We can also index into nested structures

        >>> select = Select(0, [1, 0])

        >>> x = (torch.ones(1), (torch.ones(1)*5, torch.ones(3,3)))

        >>> select(x)

        (tensor([1.]), tensor([5.]))

        >>> select = Select([1,1])

        >>> select(x)

        (tensor([[1., 1., 1.],

                [1., 1., 1.],

                [1., 1., 1.]]),)

    """

    def __init__(self, *idxs):

        super().__init__()

        self.idxs = idxs

    def forward(self, x):  # pylint: disable=arguments-differ

        """Forward method."""

        return tuple(

            _get_nested_indices(x, idx) if isinstance(idx, (list, tuple)) else x[idx]

            for idx in self.idxs

        )

def _get_nested_indices(indexable, idxs):

    if not idxs:

        return indexable

    return _get_nested_indices(indexable[idxs[0]], idxs[1:])

"""

Expose a module which recalls the latest outputs from given modules.

"""

from functools import partial

from typing import Any, Callable, Tuple

from torch import nn

def _storer(index, storage, _module, _input, output):

    storage[index] = output

def _register_hooks(modules, storage):

    return [

        mod.register_forward_hook(partial(_storer, mod, storage)) for mod in modules

    ]

def _remove_hooks(handlers):

    while handlers:

        handlers.pop().remove()

def _collect(modules, storage):

    return tuple(storage[mod] for mod in modules)

class Recall(nn.Module):

    """Recall the latest outputs from the given modules.

    Args:

        modules: ``nn.Module`` objects to recall the outputs of during this

          object's forward pass.

        agg_fn: A function used to aggregate the outputs from the supplied

          ``nn.Module`` objects. NB the input is always a single tuple

          containing the inputs into the forward, followed by the outputs

          from the latest forwards from ``modules``.

    Example:

        >>>

        # In this example we construct a simple network where the last module

        # receives input from both a module we construct first (``linear``)

        # and the previous module. The aggregation simply returns its input

        # (which will be a tuple).

        >>> import torch

        >>> linear = nn.Linear(3, 4)

        >>> net = nn.Sequential(

        ...     nn.Sequential(nn.Linear(2, 3), linear),

        ...     nn.Sequential(nn.Linear(4, 5), Recall(linear, agg_fn=lambda x: x)),

        ... )

        >>> output = net(torch.ones(10, 2))

        >>> assert isinstance(output, tuple)

        >>>

        >>> # We can also use `get_nested_modules` to retrieve outputs from a

        >>> # instanciated network:

        >>> from borch.utils.module_utils import get_nested_modules

        >>> modules = get_nested_modules(net, [("0", "0"), ("1", "0")])

        >>> recall = Recall(*modules, agg_fn=lambda x: x)

        >>> _ = net(torch.ones(10, 2))

        >>> outputs = recall()

        >>> out1 = outputs[0]

        >>> out2 = outputs[1]

    """

    # pylint: disable=abstract-method

    def __init__(self, *modules: nn.Module, agg_fn: Callable[[Tuple[Any]], Any]):

        super().__init__()

        self._agg_fn = agg_fn

        self._storage = {}

        self._collect_from = modules

        self._handlers = _register_hooks(self._collect_from, self._storage)

    def forward(self, *x):

        """Forward method."""

        return self._agg_fn(x + _collect(self._collect_from, self._storage))

    def __del__(self):

        _remove_hooks(self._handlers)