def __init__()

in optimum/tpu/modeling_llama.py [0:0]

• 29 lines of code
• 1 McCabe index (conditional complexity)

    def __init__(self, config, rank=0, world_size=1):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.rank = rank
        self.world_size = world_size
        self.gate_proj = ColumnParallelLinear.create(
            self.hidden_size,
            self.intermediate_size,
            bias=False,
            rank=rank,
            world_size=world_size,
        )
        self.up_proj = ColumnParallelLinear.create(
            self.hidden_size,
            self.intermediate_size,
            bias=False,
            rank=rank,
            world_size=world_size,
        )
        self.down_proj = RowParallelLinear.create(
            self.intermediate_size,
            self.hidden_size,
            bias=False,
            rank=rank,
            world_size=world_size,
        )
        self.act_fn = ACT2FN[config.hidden_act]