Consolidate NDTensors serialization into one file organized by type

Merges NDTensors/src/serialization_types.jl into NDTensors/src/serialization.jl
and reorders the file so each storage type's section contains:
  1. The Serialized* struct definition (with docstring).
  2. The serialized_type declaration.
  3. The Base.convert overloads in both directions.

Keeps the shared block-offset helpers near the top of the file since they're
used by two sections (BlockSparse, DiagBlockSparse). serialization_types.jl
is deleted.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: mtfishman

NDTensors/src/NDTensors.jl

@@ -90,7 +90,6 @@ include("empty/adapt.jl")
 #####################################
 # Serialization (schema types + backend-agnostic conversions, used by NDTensorsJLD2Ext)
 #
-include("serialization_types.jl")
 include("serialization.jl")
 # Mark the on-disk schema types and the `serialized_type` mapping as public API (per
 # Julia docs recommendation for backwards-compatible `public` declarations on pre-1.11

NDTensors/src/serialization.jl

@@ -1,10 +1,23 @@
 # Backend-agnostic serialization layer for NDTensors storage types.
 #
-# Pairs each storage type with its [`SerializedX`](@ref) schema struct via
-# [`serialized_type`](@ref), and defines `Base.convert` overloads to do the value-level
-# transform in both directions. Serialization backends (today only `NDTensorsJLD2Ext`)
-# layer on top: they map their own writeas mechanism to `serialized_type`, and the
-# `Base.convert` overloads here handle the actual byte-shape transform.
+# Each section below defines, for one in-memory storage type:
+#   1. The [`SerializedX`](@ref) schema struct (the on-disk layout).
+#   2. The [`serialized_type`](@ref) declaration mapping the in-memory type to the
+#      schema struct, used by serialization backends to bridge into their own writeas /
+#      type-mapping mechanism.
+#   3. `Base.convert` overloads for the value-level transform in both directions.
+#
+# JLD2's default `wconvert` / `rconvert` already delegate to `Base.convert`, so backends
+# (today only `NDTensorsJLD2Ext`) only need to register a `JLD2.writeas` declaration
+# pointing at `serialized_type` — no JLD2-specific value-conversion code is required.
+#
+# Integer-width conventions for cross-language readability:
+#   * `version::UInt32` matches `Base.VersionNumber`'s field width.
+#   * Block-sparse layouts use `block_indices::Matrix{Int64}` shaped `(ndims, num_blocks)`,
+#     following the COO sparse-tensor convention used by Apache Arrow, PyData Sparse, and
+#     PyTorch's sparse format. The tensor rank is implicit in `size(block_indices, 1)`
+#     and is preserved even when `num_blocks == 0` because HDF5 stores both matrix
+#     dimensions.
 
 """
     NDTensors.serialized_type(::Type{T}) -> Type
@@ -19,21 +32,7 @@ NDTensors.serialized_type(Dense{Float64})  # => NDTensors.SerializedDense{Float6
 """
 function serialized_type end
 
-serialized_type(::Type{<:Dense{T}}) where {T} = SerializedDense{T}
-serialized_type(::Type{<:BlockSparse{T}}) where {T} = SerializedBlockSparse{T}
-serialized_type(::Type{<:NonuniformDiag{T}}) where {T} = SerializedDiag{T}
-serialized_type(::Type{<:UniformDiag{T}}) where {T} = SerializedUniformDiag{T}
-function serialized_type(::Type{<:NonuniformDiagBlockSparse{T}}) where {T}
-    return SerializedDiagBlockSparse{T}
-end
-function serialized_type(::Type{<:UniformDiagBlockSparse{T}}) where {T}
-    return SerializedUniformDiagBlockSparse{T}
-end
-serialized_type(::Type{<:EmptyStorage{T}}) where {T} = SerializedEmptyStorage{T}
-
-# Shared block-offset (de)serialization helpers, used here for BlockSparse and
-# DiagBlockSparse. Block positions are written as the columns of a
-# `(ndims, num_blocks)` `Matrix{Int64}` (COO sparse-tensor convention).
+# Shared block-offset (de)serialization helpers, used by `BlockSparse` and `DiagBlockSparse`.
 
 function _serialize_blockoffsets(storage, ::Val{N}) where {N}
     boffs = blockoffsets(storage)
@@ -56,7 +55,24 @@ function _deserialize_blockoffsets(s, ::Val{N}) where {N}
     return boffs
 end
 
-# Dense
+# --- Dense ---
+
+"""
+    SerializedDense{T}
+
+On-disk schema for `Dense{T}` storage. Version 1.
+
+Fields:
+
+  - `version::UInt32`
+  - `data::Vector{T}`
+"""
+struct SerializedDense{T}
+    version::UInt32
+    data::Vector{T}
+end
+
+serialized_type(::Type{<:Dense{T}}) where {T} = SerializedDense{T}
 
 function Base.convert(::Type{SerializedDense{T}}, d::Dense{T}) where {T}
     return SerializedDense{T}(UInt32(1), convert(Vector{T}, data(d)))
@@ -68,7 +84,30 @@ function Base.convert(::Type{S}, s::SerializedDense) where {T, S <: Dense{T}}
     return S(s.data)
 end
 
-# BlockSparse
+# --- BlockSparse ---
+
+"""
+    SerializedBlockSparse{T}
+
+On-disk schema for `BlockSparse{T}` storage. Version 1.
+
+Fields:
+
+  - `version::UInt32`
+  - `data::Vector{T}` — flat element buffer.
+  - `block_indices::Matrix{Int64}` — shape `(ndims, num_blocks)`, each column a block
+    position (COO convention; tensor rank is `size(block_indices, 1)`).
+  - `block_offsets::Vector{Int64}` — length `num_blocks`, the offset into `data` for each
+    block.
+"""
+struct SerializedBlockSparse{T}
+    version::UInt32
+    data::Vector{T}
+    block_indices::Matrix{Int64}
+    block_offsets::Vector{Int64}
+end
+
+serialized_type(::Type{<:BlockSparse{T}}) where {T} = SerializedBlockSparse{T}
 
 function Base.convert(
         ::Type{SerializedBlockSparse{T}}, bs::BlockSparse{T, <:Any, N}
@@ -87,7 +126,24 @@ function Base.convert(::Type{S}, s::SerializedBlockSparse) where {T, S <: BlockS
     return BlockSparse(s.data, _deserialize_blockoffsets(s, Val(N)))::S
 end
 
-# Diag (nonuniform)
+# --- Diag (nonuniform) ---
+
+"""
+    SerializedDiag{T}
+
+On-disk schema for non-uniform `Diag{T}` storage. Version 1.
+
+Fields:
+
+  - `version::UInt32`
+  - `data::Vector{T}` — the diagonal entries.
+"""
+struct SerializedDiag{T}
+    version::UInt32
+    data::Vector{T}
+end
+
+serialized_type(::Type{<:NonuniformDiag{T}}) where {T} = SerializedDiag{T}
 
 function Base.convert(::Type{SerializedDiag{T}}, d::NonuniformDiag{T}) where {T}
     return SerializedDiag{T}(UInt32(1), convert(Vector{T}, data(d)))
@@ -100,7 +156,24 @@ function Base.convert(::Type{S}, s::SerializedDiag) where {T, S <: NonuniformDia
     return Diag{T}(s.data)::S
 end
 
-# Diag (uniform)
+# --- Diag (uniform) ---
+
+"""
+    SerializedUniformDiag{T}
+
+On-disk schema for uniform `Diag{T}` storage (all diagonal entries equal). Version 1.
+
+Fields:
+
+  - `version::UInt32`
+  - `value::T` — the shared diagonal value.
+"""
+struct SerializedUniformDiag{T}
+    version::UInt32
+    value::T
+end
+
+serialized_type(::Type{<:UniformDiag{T}}) where {T} = SerializedUniformDiag{T}
 
 function Base.convert(::Type{SerializedUniformDiag{T}}, d::UniformDiag{T}) where {T}
     return SerializedUniformDiag{T}(UInt32(1), convert(T, data(d)))
@@ -113,7 +186,31 @@ function Base.convert(::Type{S}, s::SerializedUniformDiag) where {T, S <: Unifor
     return Diag(s.value)::S
 end
 
-# DiagBlockSparse (nonuniform)
+# --- DiagBlockSparse (nonuniform) ---
+
+"""
+    SerializedDiagBlockSparse{T}
+
+On-disk schema for non-uniform `DiagBlockSparse{T}` storage. Version 1. Layout matches
+[`SerializedBlockSparse`](@ref).
+
+Fields:
+
+  - `version::UInt32`
+  - `data::Vector{T}`
+  - `block_indices::Matrix{Int64}` shape `(ndims, num_blocks)`
+  - `block_offsets::Vector{Int64}` length `num_blocks`
+"""
+struct SerializedDiagBlockSparse{T}
+    version::UInt32
+    data::Vector{T}
+    block_indices::Matrix{Int64}
+    block_offsets::Vector{Int64}
+end
+
+function serialized_type(::Type{<:NonuniformDiagBlockSparse{T}}) where {T}
+    return SerializedDiagBlockSparse{T}
+end
 
 function Base.convert(
         ::Type{SerializedDiagBlockSparse{T}},
@@ -135,7 +232,30 @@ function Base.convert(
     return DiagBlockSparse(s.data, _deserialize_blockoffsets(s, Val(N)))::S
 end
 
-# DiagBlockSparse (uniform)
+# --- DiagBlockSparse (uniform) ---
+
+"""
+    SerializedUniformDiagBlockSparse{T}
+
+On-disk schema for uniform `DiagBlockSparse{T}` storage. Version 1.
+
+Fields:
+
+  - `version::UInt32`
+  - `value::T` — the shared diagonal value.
+  - `block_indices::Matrix{Int64}` shape `(ndims, num_blocks)`
+  - `block_offsets::Vector{Int64}` length `num_blocks`
+"""
+struct SerializedUniformDiagBlockSparse{T}
+    version::UInt32
+    value::T
+    block_indices::Matrix{Int64}
+    block_offsets::Vector{Int64}
+end
+
+function serialized_type(::Type{<:UniformDiagBlockSparse{T}}) where {T}
+    return SerializedUniformDiagBlockSparse{T}
+end
 
 function Base.convert(
         ::Type{SerializedUniformDiagBlockSparse{T}},
@@ -157,14 +277,26 @@ function Base.convert(
     return DiagBlockSparse(s.value, _deserialize_blockoffsets(s, Val(N)))::S
 end
 
-# EmptyStorage
+# --- EmptyStorage ---
 
-function Base.convert(
-        ::Type{SerializedEmptyStorage{T}}, ::EmptyStorage{T}
-    ) where {T}
-    return SerializedEmptyStorage{T}(UInt32(1))
+"""
+    SerializedEmptyStorage{T}
+
+On-disk schema for `EmptyStorage{T}`. Version 1. The element type `T` is carried in the
+parametric type name on disk, so no data field is needed.
+
+Fields:
+
+  - `version::UInt32`
+"""
+struct SerializedEmptyStorage{T}
+    version::UInt32
 end
 
-function Base.convert(::Type{S}, ::SerializedEmptyStorage) where {T, S <: EmptyStorage{T}}
-    return S()
+serialized_type(::Type{<:EmptyStorage{T}}) where {T} = SerializedEmptyStorage{T}
+
+function Base.convert(::Type{SerializedEmptyStorage{T}}, ::EmptyStorage{T}) where {T}
+    return SerializedEmptyStorage{T}(UInt32(1))
 end
+
+Base.convert(::Type{S}, ::SerializedEmptyStorage) where {T, S <: EmptyStorage{T}} = S()