{"record":{"id":"1e0d72f66ed4d01c","repo":"Lightning-AI/pytorch-lightning","slug":"currently-the-xlastrategy-only-supports-sum-m-1e0d72","errorCode":null,"errorMessage":"Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got: {reduce_op}","messagePattern":"Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/strategies/xla.py","lineNumber":261,"sourceCode":"        else:\n            obj = obj.to(original_device)\n\n        return obj\n\n    @override\n    def reduce(\n        self,\n        output: Union[Tensor, Any],\n        group: Optional[Any] = None,\n        reduce_op: Optional[Union[ReduceOp, str]] = \"mean\",\n    ) -> Tensor:\n        if not isinstance(output, Tensor):\n            output = torch.tensor(output, device=self.root_device)\n\n        invalid_reduce_op = isinstance(reduce_op, ReduceOp) and reduce_op != ReduceOp.SUM\n        invalid_reduce_op_str = isinstance(reduce_op, str) and reduce_op.lower() not in (\"sum\", \"mean\", \"avg\")\n        if invalid_reduce_op or invalid_reduce_op_str:\n            raise ValueError(\n                \"Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got:\"\n                f\" {reduce_op}\"\n            )\n\n        import torch_xla.core.xla_model as xm\n\n        output = xm.mesh_reduce(\"reduce\", output, sum)\n\n        if isinstance(reduce_op, str) and reduce_op.lower() in (\"avg\", \"mean\"):\n            output = output / self.world_size\n\n        return output\n\n    @override\n    def setup_environment(self) -> None:\n        self._launched = True\n        super().setup_environment()\n","sourceCodeStart":243,"sourceCodeEnd":279,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/strategies/xla.py#L243-L279","documentation":"XLAStrategy.reduce only supports SUM-like reductions because it maps onto torch_xla's xm.reduce_type / all-reduce, which supports 'sum', 'mean', 'avg' (and ReduceOp.SUM). Requesting other ops (MAX, MIN, PRODUCT, etc.) raises ValueError.","triggerScenarios":"Calling strategy.reduce(tensor, reduce_op='max') or ReduceOp.MAX / any op other than sum/mean/avg; using LightningModule.all_gather/merge things or custom code invoking reduce with unsupported ops under XLAStrategy.","commonSituations":"Custom metrics that reduce max/min across ranks; porting DDP code where torch.distributed.ReduceOp.MAX works fine.","solutions":["Compute sum-based reductions instead (e.g. transform max into -min of negatives: reduce(-x, 'min-free' sum trick)) or implement max via all_gather + local max","Use 'sum'/'mean'/'avg' only when calling strategy.reduce","Do the min/max locally after an all_gather of values"],"exampleFix":"# before\nval = self.trainer.strategy.reduce(tensor, reduce_op=\"max\")  # ValueError\n\n# after\ngathered = self.trainer.strategy.all_gather(tensor)\nval = gathered.max(dim=0).values  # compute max after gather","handlingStrategy":"validation","validationCode":"op = str(reduce_op).lower()\nassert op in (\"sum\", \"mean\", \"avg\") or reduce_op in (None, ReduceOp.SUM), \"XLA reduce supports only sum/mean/avg\"","typeGuard":"def xla_reduce_ok(reduce_op) -> bool:\n    import torch.distributed as dist\n    if isinstance(reduce_op, dist.ReduceOp):\n        return reduce_op == dist.ReduceOp.SUM\n    return isinstance(reduce_op, str) and reduce_op.lower() in (\"sum\", \"mean\", \"avg\")","tryCatchPattern":null,"preventionTips":["Restrict cross-rank reductions to sum/mean/avg under XLA","Implement max/min via all_gather then local reduction","Abstract reduce ops behind an accelerator-aware helper"],"tags":["xla","reduce","reduce-op","distributed"],"backgroundTag":"unsupported-reduce-op","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}