polars.DataFrame.to_jax

DataFrame.to_jax(
    return_type: JaxExportType = 'array',
    *,
    device: jax.Device | str | None = None,
    label: str | Expr | Sequence[str | Expr] | None = None,
    features: str | Expr | Sequence[str | Expr] | None = None,
    dtype: PolarsDataType | None = None,
    order: IndexOrder = 'fortran',
) → jax.Array | dict[str, jax.Array]

Преобразовать DataFrame в массив Jax или словарь массивов Jax.

Добавлено в версии 0.20.27.

Предупреждение

В настоящее время эта функциональность считается нестабильной. Она может быть изменена в любой момент без того, чтобы это считалось нарушающим обратную совместимость изменением.

Параметры:

return_type{“array”, “dict”}

Задать тип возвращаемого значения: массив Jax или словарь массивов Jax.

device

Указать устройство Jax Device, на котором будет создан массив; можно передать строку (например, “cpu”, “gpu” или “tpu”), и в этом случае устройство будет получено как jax.devices(string)[0]. Для более точного управления можно напрямую передать экземпляр Device. Если указано None, массивы создаются на устройстве по умолчанию.

label

Одно или несколько имён столбцов, выражений или селекторов, задающих данные меток; в результате будет возвращён словарь {"label": ..., "features": ...}, если return_type имеет значение “dict”, а не словарь {"col": array, }.

features

Одно или несколько имён столбцов, выражений или селекторов, содержащих данные признаков; если параметр не указан, используются все столбцы, не отнесённые к меткам. Применяется, только если return_type имеет значение “dict”.

dtype

Привести типы всех возвращаемых массивов к единому типу; перед преобразованием в Array столбцы, тип которых отличается от требуемого, будут приведены к нужному типу. Обратите внимание: при экспорте по умолчанию используется одинарная точность (32 бита), если конфигурация или среда Jax не указывают иное (например, если при запуске в объекте конфигурации было установлено значение True для “jax_enable_x64” или в среде задано “JAX_ENABLE_X64” со значением “1”).

order{“c”, “fortran”}

Порядок индексов возвращаемого массива Jax: C-подобный (построчный) или Fortran-подобный (постолбцовый).

См. также

to_dummies

to_numpy

to_torch

Примеры

>>> df = pl.DataFrame(
...     {
...         "lbl": [0, 1, 2, 3],
...         "feat1": [1, 0, 0, 1],
...         "feat2": [1.5, -0.5, 0.0, -2.25],
...     }
... )

Стандартный тип возвращаемого значения (двумерный Array) на устройстве по умолчанию:

>>> df.to_jax()
Array([[ 0.  ,  1.  ,  1.5 ],
       [ 1.  ,  0.  , -0.5 ],
       [ 2.  ,  0.  ,  0.  ],
       [ 3.  ,  1.  , -2.25]], dtype=float32)

Создать Array на графическом процессоре по умолчанию:

>>> a = df.to_jax(device="gpu")  
>>> a.device()  
GpuDevice(id=0, process_index=0)

Создать Array на указанном графическом процессоре:

>>> gpu_device = jax.devices("gpu")[1]  
>>> a = df.to_jax(device=gpu_device)  
>>> a.device()  
GpuDevice(id=1, process_index=0)

В виде словаря отдельных массивов:

>>> df.to_jax("dict")
{'lbl': Array([0, 1, 2, 3], dtype=int32),
 'feat1': Array([1, 0, 0, 1], dtype=int32),
 'feat2': Array([ 1.5 , -0.5 ,  0.  , -2.25], dtype=float32)}

В виде словаря “label” и “features”; обратите внимание: поскольку “features” не указан, по умолчанию используются все столбцы, не входящие в “label”:

>>> df.to_jax("dict", label="lbl")
{'label': Array([[0],
        [1],
        [2],
        [3]], dtype=int32),
 'features': Array([[ 1.  ,  1.5 ],
        [ 0.  , -0.5 ],
        [ 0.  ,  0.  ],
        [ 1.  , -2.25]], dtype=float32)}

В виде словаря “label” и “features”, где для каждого параметра задано выражение col или выражение-селектор (их также можно использовать для приведения типов данных, если для меток и признаков лучше подходят разные типы):

>>> 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)}