feat(python): Add to_jax methods to support Jax Array export from `…

…DataFrame` and `Series` (#16294)
pola-rs · May 19, 2024 · 18c5601 · 18c5601
1 parent fa6b1ca
commit 18c5601
Show file tree

Hide file tree

Showing 13 changed files with 827 additions and 311 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/aggregate/mod.rs b/crates/polars-core/src/chunked_array/ops/aggregate/mod.rs
@@ -195,7 +195,7 @@ where
     }
 }
 
-/// Booleans are casted to 1 or 0.
+/// Booleans are cast to 1 or 0.
 impl BooleanChunked {
     pub fn sum(&self) -> Option<IdxSize> {
         Some(if self.is_empty() {

diff --git a/py-polars/docs/source/reference/dataframe/export.rst b/py-polars/docs/source/reference/dataframe/export.rst
@@ -13,6 +13,7 @@ Export DataFrame data to other formats:
     DataFrame.to_dict
     DataFrame.to_dicts
     DataFrame.to_init_repr
+    DataFrame.to_jax
     DataFrame.to_numpy
     DataFrame.to_pandas
     DataFrame.to_struct

diff --git a/py-polars/docs/source/reference/series/export.rst b/py-polars/docs/source/reference/series/export.rst
@@ -10,7 +10,9 @@ Export Series data to other formats:
 
    Series.to_arrow
    Series.to_frame
+   Series.to_jax
    Series.to_list
    Series.to_numpy
    Series.to_pandas
    Series.to_init_repr
+   Series.to_torch