from __future__ import annotations
from typing import TYPE_CHECKING
import paddle
from paddle.nn import Linear
import paddle.nn.functional as F
from paddle.jit.marker import unified
from paddle.jit import to_static
if TYPE_CHECKING:
    from paddle import Tensor
    from paddle._typing import (
        ParamAttrLike,
    )
@unified
class MyLinear(paddle.nn.Layer):
    weight: Tensor
    bias: Tensor
    name: str | None
    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight_attr: ParamAttrLike | None = None,
        bias_attr: ParamAttrLike | None = None,
        name: str | None = None,
    ) -> None:
        super().__init__()
        self._dtype = self._helper.get_default_dtype()
        self._weight_attr = weight_attr
        self._bias_attr = bias_attr
        self.weight = self.create_parameter(
            shape=[in_features, out_features],
            attr=self._weight_attr,
            dtype=self._dtype,
            is_bias=False,
        )
        self.bias = self.create_parameter(
            shape=[out_features],
            attr=self._bias_attr,
            dtype=self._dtype,
            is_bias=True,
        )
        self.name = name
    def forward(self, input: Tensor) -> Tensor:
        out = F.linear(
            x=input, weight=self.weight, bias=self.bias, name=self.name
        )
        return out
def call_linear(x, linear):
    return linear(x)
if __name__ == "__main__":
    linear = MyLinear(32, 10)
    static_fn = to_static(call_linear)
    x = paddle.randn([4, 32])
    y = static_fn(x, linear)
    print(type(linear.weight), linear.weight.shape)
    print(type(linear.bias), linear.weight.shape)
    print(y.shape)