pola-rs/polars · error · ValueError

`label` and `features` only apply when `return_type` is 'dic

Error message

`label` and `features` only apply when `return_type` is 'dict'

What it means

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.

Source

Thrown at py-polars/src/polars/dataframe/frame.py:2235

        >>> import polars.selectors as cs
        >>> df.to_jax(
        ...     return_type="dict",
        ...     features=cs.float(),
        ...     label=pl.col("lbl").cast(pl.UInt8),
        ... )
        {'label': Array([[0],
                [1],
                [2],
                [3]], dtype=uint8),
         'features': Array([[ 1.5 ],
                [-0.5 ],
                [ 0.  ],
                [-2.25]], dtype=float32)}
        """
        if return_type != "dict" and (label is not None or features is not None):
            msg = "`label` and `features` only apply when `return_type` is 'dict'"
            raise ValueError(msg)
        elif return_type == "dict" and label is None and features is not None:
            msg = "`label` is required if setting `features` when `return_type='dict'"
            raise ValueError(msg)

        jx = import_optional(
            "jax",
            install_message="Please see `https://jax.readthedocs.io/en/latest/installation.html` "
            "for specific installation recommendations for the Jax package",
        )
        enabled_double_precision = jx.config.jax_enable_x64 or bool(
            int(os.environ.get("JAX_ENABLE_X64", "0"))
        )
        if dtype:
            frame = self.cast(dtype)
        elif not enabled_double_precision:
            # enforce single-precision unless environment/config directs otherwise
            frame = self.cast({Float64: Float32, Int64: Int32, UInt64: UInt32})
        else:

View on GitHub (pinned to df599052da)

Solutions

1. Set return_type='dict' when using label/features: `df.to_jax('dict', label='y', features=['x'])`
2. Or drop label/features and export the full frame: `df.to_jax()`
3. Select columns explicitly beforehand if you only need features: `df.select(features).to_jax()`

Example fix

# before
train = df.to_jax(label='target', features=['f1', 'f2'])

# after
train = df.to_jax('dict', label='target', features=['f1', 'f2'])

if (label is not None or features is not None) and return_type != 'dict':
    raise ValueError('to_jax: label/features require return_type="dict"')
out = df.to_jax(return_type, label=label, features=features)

try:
    out = df.to_jax(return_type, label=label, features=features)
except ValueError as e:
    if 'only apply when' in str(e):
        out = df.to_jax('dict', label=label, features=features)
    else:
        raise

When it happens

Trigger: `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.

Common situations: 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.

Related errors

• `label` is required if setting `features` when `return_type=
• invalid `return_type`: {return_type!r} Expected one of: {val
• `label` and `features` only apply when `return_type` is 'dat
• invalid `return_type`: {return_type!r} Expected one of: {val
• cannot call `.item()` with only one of `row` or `column`

AI-assisted analysis of pola-rs/polars@df599052da (2026-08-16). Data as JSON: /api/errors/367b64561740577a. Report an issue: GitHub.