jax authors cd7f03f272 Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays. ...

Currently the serialized data and its length are being carried in two separate arrays, a fixed-with bytes array (with a hard-coded max size) and a unit32 array respectively.

PiperOrigin-RevId: 734299259

2025-03-06 14:57:52 -08:00

..

__init__.py

Update references to JAX's GitHub repo

2025-03-04 18:31:09 -08:00

api.py

[JAX] Implement an initial object API for colocated Python

2025-02-21 12:58:25 -08:00

func_backend.py

[JAX] Add the function API of jax.experimental.colocated_python

2024-10-28 12:18:48 -07:00

func.py

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.

2025-03-06 14:57:52 -08:00

obj_backend.py

[JAX] Implement an initial object API for colocated Python

2025-02-21 12:58:25 -08:00

obj.py

[JAX] Implement an initial object API for colocated Python

2025-02-21 12:58:25 -08:00

serialization.py

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.

2025-03-06 14:57:52 -08:00