`sub_byte_element_size_in_bits` is a lowering only thing for now (since we know the dtype of the aval so JAX can add the appropriate value). We can expose it to the user API if required.
memory space is exposed via JAX memories API so it doesn't have to be in the layout API.
Also expose `_xla_layout` as a private API from `PJRTLayout` so that we can access fields to create JAX layouts.
Add construtors to `xla::Layout` so that JAX can create Layouts with minor_to_major and tiling information.
PiperOrigin-RevId: 647487510
Previously, when writing the OCDBT format, the manifest and root B+tree node could be redundantly written multiple times depending on timing.
With this change, the manifest and root B+tree node are always written only once.
Additionally, source data was previously redundantly copied into the TensorStore chunk cache.
PiperOrigin-RevId: 642345928
PjRtCApiBuffer::layout() was using PJRT_Buffer_GetMemoryLayout, which will be deprecated. PJRT_Buffer_GetMemoryLayout uses explicit PJRT_Buffer_MemoryLayout which does not contain all the layout information.
PiperOrigin-RevId: 638048293
Before, dtype used to be in the metadata field of tensorstore spec because of it was the legacy way to config the dtype. This setting doesn't understand the "str" name, hence, there was special logic to translate bfloat for example.
This CL moves it out of the metadata field and put the dtype directly into the Tensorstore spec to eliminate special dtype translation logic. This will also add support of other quantized types such as int4.
PiperOrigin-RevId: 629845048
This makes in-flight memory limiter both reflective of the actual peak usage, as well as reduces peak usage since we no longer try to fully materialize sharded tensors on the host.
PiperOrigin-RevId: 524456216
