def forward()

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

• 18 lines of code
• 4 McCabe index (conditional complexity)

    def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
        input_parallel = copy_to_model_parallel_region(input_, self.groups, self.world_size, self.rank)
        # PyTorch eager and inductor do not accept negative values in the input to embedding
        # layers. Take the modulus to avoid this error.
        if USE_CUDA:
            input_parallel = torch.remainder(input_parallel, self.weight.shape[0])
        weight = self.weight
        if self.quant:
            weight = weight * self.weight_scaler.unsqueeze(-1)
        output_parallel = F.embedding(
            input_parallel,
            weight,
            self.padding_idx,
            self.max_norm,
            self.norm_type,
            self.scale_grad_by_freq,
            self.sparse,
        )
        output = gather_from_model_parallel_region(output_parallel, self.groups, self.world_size, self.rank)
        return output