{"record":{"id":"367b64561740577a","repo":"pola-rs/polars","slug":"label-and-features-only-apply-when-return-typ","errorCode":null,"errorMessage":"`label` and `features` only apply when `return_type` is 'dict'","messagePattern":"`label` and `features` only apply when `return_type` is 'dict'","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"py-polars/src/polars/dataframe/frame.py","lineNumber":2235,"sourceCode":"\n        >>> import polars.selectors as cs\n        >>> df.to_jax(\n        ...     return_type=\"dict\",\n        ...     features=cs.float(),\n        ...     label=pl.col(\"lbl\").cast(pl.UInt8),\n        ... )\n        {'label': Array([[0],\n                [1],\n                [2],\n                [3]], dtype=uint8),\n         'features': Array([[ 1.5 ],\n                [-0.5 ],\n                [ 0.  ],\n                [-2.25]], dtype=float32)}\n        \"\"\"\n        if return_type != \"dict\" and (label is not None or features is not None):\n            msg = \"`label` and `features` only apply when `return_type` is 'dict'\"\n            raise ValueError(msg)\n        elif return_type == \"dict\" and label is None and features is not None:\n            msg = \"`label` is required if setting `features` when `return_type='dict'\"\n            raise ValueError(msg)\n\n        jx = import_optional(\n            \"jax\",\n            install_message=\"Please see `https://jax.readthedocs.io/en/latest/installation.html` \"\n            \"for specific installation recommendations for the Jax package\",\n        )\n        enabled_double_precision = jx.config.jax_enable_x64 or bool(\n            int(os.environ.get(\"JAX_ENABLE_X64\", \"0\"))\n        )\n        if dtype:\n            frame = self.cast(dtype)\n        elif not enabled_double_precision:\n            # enforce single-precision unless environment/config directs otherwise\n            frame = self.cast({Float64: Float32, Int64: Int32, UInt64: UInt32})\n        else:","sourceCodeStart":2217,"sourceCodeEnd":2253,"githubUrl":"https://github.com/pola-rs/polars/blob/df599052daf96e7a9cc30a3b0c6bd25d6947e3c0/py-polars/src/polars/dataframe/frame.py#L2217-L2253","documentation":"Raised by DataFrame.to_jax when `label` or `features` is supplied but `return_type` is not 'dict'. Label/features splitting only has a defined meaning for the dict export, which returns {'label': array, 'features': array}; for the default 'array' return type there is nowhere to put a separate label tensor. Polars validates this combination before importing jax, so it fires even without jax installed.","triggerScenarios":"`df.to_jax(label='y')`, `df.to_jax('array', features=['x1','x2'])`, or any to_jax call with return_type other than 'dict' (or the default) while label/features is not None.","commonSituations":"Refactoring model-prep code and dropping the `\"dict\"` first argument; sharing a helper between to_torch('dataset', label=...) and to_jax(...) and forgetting jax's dict requirement; upgrading from older polars where the call shape differed.","solutions":["Set return_type='dict' when using label/features: `df.to_jax('dict', label='y', features=['x'])`","Or drop label/features and export the full frame: `df.to_jax()`","Select columns explicitly beforehand if you only need features: `df.select(features).to_jax()`"],"exampleFix":"# before\ntrain = df.to_jax(label='target', features=['f1', 'f2'])\n\n# after\ntrain = df.to_jax('dict', label='target', features=['f1', 'f2'])","handlingStrategy":"validation","validationCode":"if (label is not None or features is not None) and return_type != 'dict':\n    raise ValueError('to_jax: label/features require return_type=\"dict\"')\nout = df.to_jax(return_type, label=label, features=features)","typeGuard":null,"tryCatchPattern":"try:\n    out = df.to_jax(return_type, label=label, features=features)\nexcept ValueError as e:\n    if 'only apply when' in str(e):\n        out = df.to_jax('dict', label=label, features=features)\n    else:\n        raise","preventionTips":["Pair label/features with return_type='dict' in every to_jax call","Build export kwargs once and validate the (return_type, label, features) triple","Remember to_jax valid types are only 'array' and 'dict'"],"tags":["jax","export","argument-mismatch","ml"],"backgroundTag":null,"analyzedSha":"df599052daf96e7a9cc30a3b0c6bd25d6947e3c0","analyzedAt":"2026-08-16T12:10:03.978Z","schemaVersion":2},"datasetVersion":"2026-08-16T13:17:31.715Z"}