File size: 1,401 Bytes
d53adc9 | 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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | # Generated by ml.integrations.export.runtime_packager.write_remote_code_bundle.
# Exported for HuggingFace trust_remote_code loading.
# This file is intentionally self-contained.
"""Runtime linear layers used by the model."""
from __future__ import annotations
import torch
import torch.nn.functional as functional
from torch import nn
class RuntimeLinear(nn.Module):
"""Plain linear layer used across the runtime model."""
def __init__(
self,
in_features: int,
out_features: int,
*,
bias: bool = False,
) -> None:
super().__init__()
self.in_features = int(in_features)
self.out_features = int(out_features)
self.weight = nn.Parameter(torch.empty(self.out_features, self.in_features))
self.bias = nn.Parameter(torch.empty(self.out_features)) if bool(bias) else None
self.reset_parameters()
def reset_parameters(self) -> None:
nn.init.normal_(self.weight, mean=0.0, std=0.02)
if self.bias is not None:
nn.init.zeros_(self.bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return functional.linear(x, self.weight, self.bias)
def extra_repr(self) -> str:
return (
f"in_features={self.in_features}, out_features={self.out_features}, "
f"bias={self.bias is not None}"
)
__all__ = ["RuntimeLinear"]
|