Source code for k1lib.knn

# AUTOGENERATED FILE! PLEASE DON'T EDIT
"""
Some nice utils to complement :mod:`torch.nn`. This is exposed automatically
with::

   from k1lib.imports import *
   knn.Lambda # exposed
"""
from torch import nn
from typing import Callable, Any
__all__ = ["Lambda", "Identity", "LinBlock"]
class Lambda(nn.Module):
    def __init__(self, f:Callable[[Any], Any]):
        """Creates a simple module with a specified :meth:`forward`
function."""
        super().__init__(); self.f = f
    def forward(self, x): return self.f(x)
class Identity(Lambda):
    """Creates a module that returns the input in :meth:`forward`"""
    def __init__(self): super().__init__(lambda x: x)
class LinBlock(nn.Module):
    def __init__(self, inC, outC):
        """Linear layer with relu behind it"""
        super().__init__(); self.lin = nn.Linear(inC, outC); self.relu = nn.ReLU()
    def forward(self, x):
        return x | self.lin | self.relu