{"record":{"id":"60c64382c7af521f","repo":"jax-ml/jax","slug":"invalid-pipeline-stage","errorCode":null,"errorMessage":"Invalid pipeline stage","messagePattern":"Invalid pipeline stage","errorType":"validation","errorClass":"runtime_error","httpStatus":null,"severity":"error","filePath":"jaxlib/xla.cc","lineNumber":187,"sourceCode":"}  // namespace\n\nNB_MODULE(_xla, m) {\n  // Register a transform directly via RegisterHloXlaTransform.\n  // This is used for in-process backends (e.g. CPU).\n  m.def(\n      \"register_xla_transform\",\n      [](std::string name, int stage, nb::object callback) {\n        xla::HloXlaTransform::PipelineStage pipeline_stage;\n        switch (stage) {\n          case 0:\n            pipeline_stage = xla::HloXlaTransform::PipelineStage::kPreScheduler;\n            break;\n          case 1:\n            pipeline_stage =\n                xla::HloXlaTransform::PipelineStage::kPostScheduler;\n            break;\n          default:\n            throw std::runtime_error(\"Invalid pipeline stage\");\n        }\n\n        auto transform = std::make_shared<PyHloXlaTransform>(\n            std::move(name), std::move(callback));\n        xla::RegisterHloXlaTransform(pipeline_stage, std::move(transform));\n      },\n      nb::arg(\"name\"), nb::arg(\"stage\"), nb::arg(\"callback\"));\n\n  m.def(\n      \"clear_xla_transform\",\n      [](std::string name, int stage) {\n        xla::HloXlaTransform::PipelineStage pipeline_stage;\n        switch (stage) {\n          case 0:\n            pipeline_stage = xla::HloXlaTransform::PipelineStage::kPreScheduler;\n            break;\n          case 1:\n            pipeline_stage =","sourceCodeStart":169,"sourceCodeEnd":205,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/xla.cc#L169-L205","documentation":"register_xla_transform(name, callback, stage) maps stage to a PipelineStage enum (0 = pre-scheduler, 1 = post-scheduler). Any integer other than 0 or 1 throws 'Invalid pipeline stage'.","triggerScenarios":"Calling jaxlib.xla_extension.register_xla_transform with stage outside {0,1}, e.g. 2, -1, or an unmarshalled enum value.","commonSituations":"Passing a string stage name instead of an int; using a stage constant from a different JAX version; plugin/backend tooling written against a changed enum.","solutions":["Use stage 0 (kPreScheduler) or 1 (kPostScheduler) only","Prefer higher-level JAX APIs for staging rather than raw xla_extension registration"],"exampleFix":"# before\nregister_xla_transform('my_pass', cb, stage=2)\n\n# after\nregister_xla_transform('my_pass', cb, stage=1)  # 0=pre-scheduler, 1=post-scheduler","handlingStrategy":"validation","validationCode":"if stage not in (0, 1):\n    raise ValueError('stage must be 0 (pre-scheduler) or 1 (post-scheduler)')\nregister_xla_transform(name, cb, stage)","typeGuard":"def is_valid_stage(s) -> bool:\n    return s in (0, 1)","tryCatchPattern":null,"preventionTips":["Define named constants PRE_SCHEDULER=0, POST_SCHEDULER=1 instead of magic numbers"],"tags":["jax","xla","compiler-passes","validation"],"backgroundTag":"invalid-enum-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}