pola-rs/polars · error · ValueError

invalid `return_type`: {return_type!r} Expected one of: {val

Error message

invalid `return_type`: {return_type!r}
Expected one of: {valid_jax_types}

What it means

Raised by DataFrame.to_jax when `return_type` is not one of the allowed literals ('array', 'dict'). The parameter is typed as JaxExportType = Literal['array', 'dict'], and polars enumerates the valid values via get_args in the error message. Because Python does not enforce Literal at runtime, an invalid string from config or a typo reaches the dispatch chain and falls into the final else branch.

Source

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

                if label is not None:
                    # return a {"label": array(s), "features": array(s)} dict
                    label_frame = frame.select(label)
                    features_frame = (
                        frame.select(features)
                        if features is not None
                        else frame.drop(*label_frame.columns)
                    )
                    return {
                        "label": label_frame.to_jax(),
                        "features": features_frame.to_jax(),
                    }
                else:
                    # return a {"col": array} dict
                    return {srs.name: srs.to_jax() for srs in frame}
            else:
                valid_jax_types = ", ".join(get_args(JaxExportType))
                msg = f"invalid `return_type`: {return_type!r}\nExpected one of: {valid_jax_types}"
                raise ValueError(msg)

    @overload
    def to_torch(
        self,
        return_type: Literal["tensor"] = ...,
        *,
        label: str | Expr | Sequence[str | Expr] | None = ...,
        features: str | Expr | Sequence[str | Expr] | None = ...,
        dtype: PolarsDataType | None = ...,
    ) -> torch.Tensor: ...

    @overload
    def to_torch(
        self,
        return_type: Literal["dataset"],
        *,
        label: str | Expr | Sequence[str | Expr] | None = ...,
        features: str | Expr | Sequence[str | Expr] | None = ...,

View on GitHub (pinned to df599052da)

Solutions

1. Use 'array' (default) or 'dict': `df.to_jax('dict')`
2. Validate config-driven values before the call: `assert rt in ('array', 'dict')`
3. If you wanted a torch tensor/dataset, use `df.to_torch(...)` instead

Example fix

# before
arr = df.to_jax('tensor')

# after
arr = df.to_jax('array')
# or, for torch:
tensor = df.to_torch('tensor')

VALID_JAX = ('array', 'dict')
if return_type not in VALID_JAX:
    raise ValueError(f'return_type must be one of {VALID_JAX}, got {return_type!r}')
out = df.to_jax(return_type)

def is_jax_return_type(rt: object) -> bool:
    """to_jax only accepts 'array' or 'dict'."""
    return rt in ('array', 'dict')

try:
    out = df.to_jax(return_type)
except ValueError as e:
    if 'invalid `return_type`' in str(e):
        out = df.to_jax('array')
    else:
        raise

When it happens

Trigger: `df.to_jax('tensor')` (torch vocabulary), `df.to_jax('Array')`, `df.to_jax('dataset')`, or a return_type read from a config file/CLI arg that isn't exactly 'array' or 'dict'.

Common situations: Copy-pasting a to_torch return_type into a to_jax call; user-configurable export functions where the string comes from YAML/JSON; case mismatches like 'Dict'; shared constants defined for one backend and reused for another.

Related errors

• `label` and `features` only apply when `return_type` is 'dic
• `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
• `orient` must be one of {'col', 'row', None}, got {orient!r}

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