File size: 696 Bytes
1e114b1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
"""Internal helper for parameterless modules bound to the one source layer."""

from __future__ import annotations

import weakref

from torch import nn

from .source import DendroSourceLayer


class SourceBoundModule(nn.Module):
    """An ``nn.Module`` that references, but never registers, the source module."""

    def __init__(self, source: DendroSourceLayer) -> None:
        super().__init__()
        self.__dict__["_source_ref"] = weakref.ref(source)

    @property
    def source(self) -> DendroSourceLayer:
        source = self.__dict__["_source_ref"]()
        if source is None:
            raise RuntimeError("The owning DendroSourceLayer has been released")
        return source