{"record":{"id":"9f36cd8f53dec589","repo":"jax-ml/jax","slug":"output-shapes-must-match-got-and","errorCode":null,"errorMessage":"{}() output shapes must match {}, got {} and {}","messagePattern":"(.+?)\\(\\) output shapes must match (.+?), got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/solves.py","lineNumber":205,"sourceCode":"\n\ndef _transpose_one_output(linear_fun, primals):\n  transpose_fun = api.linear_transpose(linear_fun, primals)\n  def transposed_fun(x):\n    (y,) = transpose_fun(x)\n    return y\n  return transposed_fun\n\n\ndef _flatten(args):\n  return [x for arg in args for x in arg]\n\n\ndef _check_shapes(func_name, expected_name, actual, expected):\n  actual_shapes = _map(np.shape, actual)\n  expected_shapes = _map(np.shape, expected)\n  if actual_shapes != expected_shapes:\n    raise ValueError(\n        f\"{func_name}() output shapes must match {expected_name}, \"\n        f\"got {actual_shapes} and {expected_shapes}\")\n\n\n@partial(api_boundary, repro_api_name=\"jax.custom_linear_solve\")\ndef custom_linear_solve(\n    matvec: Callable,\n    b: Any,\n    solve: Callable[[Callable, Any], Any],\n    transpose_solve: Callable[[Callable, Any], Any] | None = None,\n    symmetric=False, has_aux=False):\n  \"\"\"Perform a matrix-free linear solve with implicitly defined gradients.\n\n  This function allows for overriding or defining gradients for a linear\n  solve directly via implicit differentiation at the solution, rather than by\n  differentiating *through* the solve operation. This can sometimes be much faster\n  or more numerically stable, or differentiating through the solve operation\n  may not even be implemented (e.g., if ``solve`` uses ``lax.while_loop``).","sourceCodeStart":187,"sourceCodeEnd":223,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/solves.py#L187-L223","documentation":"custom_linear_solve (and custom_root) validate that the user-supplied solve function returns arrays whose shapes exactly match the shape of the expected solution (e.g. the b matrix). The helper _check_shapes compares actual vs expected shape lists and raises ValueError showing both when they differ.","triggerScenarios":"Passing a `solve`/`solve_transpose` function to jax.lax.custom_linear_solve that returns a reshaped or transposed result, or the wrong number of arrays, e.g. returning b.T shape or a scalar instead of a vector.","commonSituations":"Custom iterative/Cholesky solvers returning solutions with extra/missing batch dims; a matvec closure that accidentally squeezes the batch dimension; mismatch between the number of outputs and the number of b arrays in pytree input.","solutions":["Make the solve function return exactly the same shapes (and pytree structure) as b","Add assert statements inside solve: lambda A, b: (assert matching shapes, x)","Return x.reshape(b.shape) as a defensive last line of the solve closure"],"exampleFix":"// before\ndef solve(A, b):\n  return jsp.linalg.solve(A, b).squeeze()  # wrong shape for b of shape (n, 1)\n// after\ndef solve(A, b):\n  return jsp.linalg.solve(A, b).reshape(b.shape)","handlingStrategy":"validation","validationCode":"x = solve(A, b)\nactual = [np.shape(a) for a in jax.tree_util.tree_leaves(x)]\nexpected = [np.shape(e) for e in jax.tree_util.tree_leaves(b)]\nassert actual == expected, (actual, expected)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["End solve closures with return x.reshape(b.shape)","Return the same pytree structure as b"],"tags":["jax","custom-linear-solve","shape-mismatch","linear-algebra"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}