Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/plugins/training_type/tpu_spawn.py (Python)

Code before:
        decision = bool(decision == self.world_size)
        return decision

    def reduce(self, output, group: Optional[Any] = None, reduce_op: Optional[Union[ReduceOp, str]] = None):
        if not isinstance(output, torch.Tensor):
            output = torch.tensor(output, device=self.lightning_module.device)

        _invalid_reduce_op = isinstance(reduce_op, ReduceOp) and reduce_op != ReduceOp.SUM
        _invalid_reduce_op_str = isinstance(reduce_op, str) and reduce_op.lower() not in ("sum", "mean", "avg")
        if _invalid_reduce_op or _invalid_reduce_op_str:
            raise MisconfigurationException(
                "Currently, TPUSpawn TrainingTypePlugin only support `sum`, `mean`, `avg` reduce operation."
            )

        output = xm.mesh_reduce("reduce", output, sum)

        if isinstance(reduce_op, str) and reduce_op.lower() in ("avg", "mean"):
            output = output / self.world_size

        return output

    def _close_logger(self, trainer) -> None:
        if trainer.logger is not None:
            trainer.logger.finalize("success")

    def get_xmp_spawn_kwargs(self, trainer: "pl.Trainer") -> dict:
        return {
            "args": (trainer, self.mp_queue),
            "nprocs": len(self.parallel_devices),
            "start_method": self.start_method,
        }

    def start_training(self, trainer: "pl.Trainer") -> None:
        # todo: precision pluging is call in accelerator setup and should be moved
        if "XLA_USE_BF16" in os.environ:
            del os.environ["XLA_USE_BF16"]
        self._close_logger(trainer)
        xmp.spawn(self.new_process, **self.get_xmp_spawn_kwargs(trainer))

    def start_evaluating(self, trainer: "pl.Trainer") -> None:
        self._close_logger(trainer)
        xmp.spawn(self.new_process, **self.get_xmp_spawn_kwargs(trainer))

    def start_predicting(self, trainer: "pl.Trainer") -> None:
        xmp.spawn(self.new_process, **self.get_xmp_spawn_kwargs(trainer))

    def training_step(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def validation_step(self, *args, **kwargs):
        return self.model(*args, **kwargs)

Review comment (suggestion):
something to address in the future: we are inheriting get_mp_spawn_args from the super class. 
I believe they should either be called the same or be made protected

Code after:
        decision = bool(decision == self.world_size)
        return decision

    def reduce(self, output, group: Optional[Any] = None, reduce_op: Optional[Union[ReduceOp, str]] = None):
        if not isinstance(output, torch.Tensor):
            output = torch.tensor(output, device=self.lightning_module.device)

        _invalid_reduce_op = isinstance(reduce_op, ReduceOp) and reduce_op != ReduceOp.SUM
        _invalid_reduce_op_str = isinstance(reduce_op, str) and reduce_op.lower() not in ("sum", "mean", "avg")
        if _invalid_reduce_op or _invalid_reduce_op_str:
            raise MisconfigurationException(
                "Currently, TPUSpawn TrainingTypePlugin only support `sum`, `mean`, `avg` reduce operation."
            )

        output = xm.mesh_reduce("reduce", output, sum)

        if isinstance(reduce_op, str) and reduce_op.lower() in ("avg", "mean"):
            output = output / self.world_size

        return output

    def _close_logger(self, trainer) -> None:
        if trainer.logger is not None:
            trainer.logger.finalize("success")

    def get_mp_spawn_kwargs(self, trainer: "pl.Trainer") -> dict:
        return {
            "args": (trainer, self.mp_queue),
            "nprocs": len(self.parallel_devices),
            "start_method": self.start_method,
        }

    def start_training(self, trainer: "pl.Trainer") -> None:
        # todo: precision pluging is call in accelerator setup and should be moved
        if "XLA_USE_BF16" in os.environ:
            del os.environ["XLA_USE_BF16"]
        self._close_logger(trainer)
        xmp.spawn(self.new_process, **self.get_mp_spawn_kwargs(trainer))

    def start_evaluating(self, trainer: "pl.Trainer") -> None:
        self._close_logger(trainer)
        xmp.spawn(self.new_process, **self.get_mp_spawn_kwargs(trainer))

    def start_predicting(self, trainer: "pl.Trainer") -> None:
        xmp.spawn(self.new_process, **self.get_mp_spawn_kwargs(trainer))

    def training_step(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def validation_step(self, *args, **kwargs):
        return self.model(*args, **kwargs)