Adds a minimal but viable implementation of string arrays (with `numpy.dtypes.StringDType`) in JAX. Currently this only supports making of a string array by means of either `jax.numpy.asarray` or `jax.device_put` and reading it back with `jax.device_get`.#21503

Merged

copybara-service[bot] merged 1 commit intomainfrom test_716042460

Feb 5, 2025

+173-5

Commits on Feb 5, 2025

Adds a minimal but viable implementation of string arrays (with numpy.dtypes.StringDType) in JAX. Currently this only supports making of a string array by means of either jax.numpy.asarray or `…
Google-ML-Automation
committed