(
self,
in_features: int,
out_features: int,
bias: bool = True,
device=None,
dtype=None,
)
| 93 | """ |
| 94 | |
| 95 | def __init__( |
| 96 | self, |
| 97 | in_features: int, |
| 98 | out_features: int, |
| 99 | bias: bool = True, |
| 100 | device=None, |
| 101 | dtype=None, |
| 102 | ) -> None: |
| 103 | super().__init__() |
| 104 | self.in_features = in_features |
| 105 | self.out_features = out_features |
| 106 | self.weight = flow.nn.Parameter( |
| 107 | flow.Tensor(out_features, in_features).to(dtype=dtype, device=device) |
| 108 | ) |
| 109 | self.bias = ( |
| 110 | flow.nn.Parameter(flow.Tensor(out_features).to(dtype=dtype, device=device)) |
| 111 | if bias |
| 112 | else None |
| 113 | ) |
| 114 | self.use_fused_matmul_bias = ( |
| 115 | self.bias is not None |
| 116 | and os.getenv("ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR") == "1" |
| 117 | ) |
| 118 | self.reset_parameters() |
| 119 | |
| 120 | def reset_parameters(self) -> None: |
| 121 | if os.getenv("ONEFLOW_LINEAR_EMBEDDING_SKIP_INIT", "0") == "1": |
no test coverage detected