Add numcodecs entry point (#21)

* Add numcodecs entry points

* Add codec encoding format versions

* Bump codec versions

* Require numcodecs-wasm v0.1.5 in numcodecs-wasm-*

Author: juntyr

Cargo.toml

@@ -42,8 +42,8 @@ rust-version = "1.82"
 
 [workspace.dependencies]
 # workspace-internal numcodecs crates
-numcodecs = { version = "0.2", path = "crates/numcodecs", default-features = false }
-numcodecs-python = { version = "0.5", path = "crates/numcodecs-python", default-features = false }
+numcodecs = { version = "0.2.1", path = "crates/numcodecs", default-features = false }
+numcodecs-python = { version = "0.6", path = "crates/numcodecs-python", default-features = false }
 numcodecs-wasm-builder = { version = "0.1", path = "crates/numcodecs-wasm-builder", default-features = false }
 numcodecs-wasm-logging = { version = "0.1", path = "crates/numcodecs-wasm-logging", default-features = false }
 numcodecs-wasm-guest = { version = "0.2", path = "crates/numcodecs-wasm-guest", default-features = false }
@@ -53,26 +53,26 @@ numcodecs-wasm-host-reproducible = { version = "0.1", path = "crates/numcodecs-w
 numcodecs-wasm = { version = "0.1", path = "py/numcodecs-wasm", default-features = false }
 
 # workspace-internal codecs crates
-numcodecs-asinh = { version = "0.2", path = "codecs/asinh", default-features = false }
-numcodecs-bit-round = { version = "0.2", path = "codecs/bit-round", default-features = false }
-numcodecs-fixed-offset-scale = { version = "0.2", path = "codecs/fixed-offset-scale", default-features = false }
-numcodecs-fourier-network = { version = "0.1", path = "codecs/fourier-network", default-features = false }
-numcodecs-identity = { version = "0.2", path = "codecs/identity", default-features = false }
-numcodecs-jpeg2000 = { version = "0.1", path = "codecs/jpeg2000", default-features = false }
-numcodecs-linear-quantize = { version = "0.3", path = "codecs/linear-quantize", default-features = false }
-numcodecs-log = { version = "0.3", path = "codecs/log", default-features = false }
-numcodecs-pco = { version = "0.1", path = "codecs/pco", default-features = false }
-numcodecs-random-projection = { version = "0.2", path = "codecs/random-projection", default-features = false }
-numcodecs-reinterpret = { version = "0.2", path = "codecs/reinterpret", default-features = false }
-numcodecs-round = { version = "0.2", path = "codecs/round", default-features = false }
-numcodecs-swizzle-reshape = { version = "0.2", path = "codecs/swizzle-reshape", default-features = false }
-numcodecs-sz3 = { version = "0.5", path = "codecs/sz3", default-features = false }
-numcodecs-tthresh = { version = "0.1", path = "codecs/tthresh", default-features = false }
-numcodecs-uniform-noise = { version = "0.2", path = "codecs/uniform-noise", default-features = false }
-numcodecs-zfp = { version = "0.4", path = "codecs/zfp", default-features = false }
-numcodecs-zfp-classic = { version = "0.2", path = "codecs/zfp-classic", default-features = false }
-numcodecs-zlib = { version = "0.2", path = "codecs/zlib", default-features = false }
-numcodecs-zstd = { version = "0.2", path = "codecs/zstd", default-features = false }
+numcodecs-asinh = { version = "0.3", path = "codecs/asinh", default-features = false }
+numcodecs-bit-round = { version = "0.3", path = "codecs/bit-round", default-features = false }
+numcodecs-fixed-offset-scale = { version = "0.3", path = "codecs/fixed-offset-scale", default-features = false }
+numcodecs-fourier-network = { version = "0.2", path = "codecs/fourier-network", default-features = false }
+numcodecs-identity = { version = "0.3", path = "codecs/identity", default-features = false }
+numcodecs-jpeg2000 = { version = "0.2", path = "codecs/jpeg2000", default-features = false }
+numcodecs-linear-quantize = { version = "0.4", path = "codecs/linear-quantize", default-features = false }
+numcodecs-log = { version = "0.4", path = "codecs/log", default-features = false }
+numcodecs-pco = { version = "0.2", path = "codecs/pco", default-features = false }
+numcodecs-random-projection = { version = "0.3", path = "codecs/random-projection", default-features = false }
+numcodecs-reinterpret = { version = "0.3", path = "codecs/reinterpret", default-features = false }
+numcodecs-round = { version = "0.3", path = "codecs/round", default-features = false }
+numcodecs-swizzle-reshape = { version = "0.3", path = "codecs/swizzle-reshape", default-features = false }
+numcodecs-sz3 = { version = "0.6", path = "codecs/sz3", default-features = false }
+numcodecs-tthresh = { version = "0.2", path = "codecs/tthresh", default-features = false }
+numcodecs-uniform-noise = { version = "0.3", path = "codecs/uniform-noise", default-features = false }
+numcodecs-zfp = { version = "0.5", path = "codecs/zfp", default-features = false }
+numcodecs-zfp-classic = { version = "0.3", path = "codecs/zfp-classic", default-features = false }
+numcodecs-zlib = { version = "0.3", path = "codecs/zlib", default-features = false }
+numcodecs-zstd = { version = "0.3", path = "codecs/zstd", default-features = false }
 
 # crates.io third-party dependencies
 anyhow = { version = "1.0.93", default-features = false }

codecs/asinh/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "numcodecs-asinh"
-version = "0.2.1"
+version = "0.3.0"
 edition = { workspace = true }
 authors = { workspace = true }
 repository = { workspace = true }

codecs/asinh/src/lib.rs

@@ -21,7 +21,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension, Zip};
 use num_traits::{Float, Signed};
 use numcodecs::{
     AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
-    Codec, StaticCodec, StaticCodecConfig,
+    Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
 };
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -46,7 +46,10 @@ use thiserror::Error;
 pub struct AsinhCodec {
     /// The width of the close-to-zero input value range where the transform is
     /// nearly linear
-    linear_width: f64,
+    pub linear_width: f64,
+    /// The codec's encoding format version. Do not provide this parameter explicitly.
+    #[serde(default, rename = "_version")]
+    pub version: StaticCodecVersion<1, 0, 0>,
 }
 
 impl Codec for AsinhCodec {
@@ -99,7 +102,7 @@ impl Codec for AsinhCodec {
 }
 
 impl StaticCodec for AsinhCodec {
-    const CODEC_ID: &'static str = "asinh";
+    const CODEC_ID: &'static str = "asinh.rs";
 
     type Config<'de> = Self;
 

codecs/bit-round/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "numcodecs-bit-round"
-version = "0.2.0"
+version = "0.3.0"
 edition = { workspace = true }
 authors = { workspace = true }
 repository = { workspace = true }

codecs/bit-round/src/lib.rs

@@ -20,7 +20,7 @@
 use ndarray::{Array, ArrayBase, Data, Dimension};
 use numcodecs::{
     AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
-    Codec, StaticCodec, StaticCodecConfig,
+    Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
 };
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -45,6 +45,9 @@ pub struct BitRoundCodec {
     /// If keepbits is equal to the bitlength of the dtype's mantissa, no
     /// transformation is performed.
     pub keepbits: u8,
+    /// The codec's encoding format version. Do not provide this parameter explicitly.
+    #[serde(default, rename = "_version")]
+    pub version: StaticCodecVersion<1, 0, 0>,
 }
 
 impl Codec for BitRoundCodec {
@@ -80,7 +83,7 @@ impl Codec for BitRoundCodec {
 }
 
 impl StaticCodec for BitRoundCodec {
-    const CODEC_ID: &'static str = "bit-round";
+    const CODEC_ID: &'static str = "bit-round.rs";
 
     type Config<'de> = Self;
 

codecs/fixed-offset-scale/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "numcodecs-fixed-offset-scale"
-version = "0.2.1"
+version = "0.3.0"
 edition = { workspace = true }
 authors = { workspace = true }
 repository = { workspace = true }

codecs/fixed-offset-scale/src/lib.rs

@@ -21,7 +21,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension};
 use num_traits::Float;
 use numcodecs::{
     AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
-    Codec, StaticCodec, StaticCodecConfig,
+    Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
 };
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -43,6 +43,9 @@ pub struct FixedOffsetScaleCodec {
     pub offset: f64,
     /// The scale of the data.
     pub scale: f64,
+    /// The codec's encoding format version. Do not provide this parameter explicitly.
+    #[serde(default, rename = "_version")]
+    pub version: StaticCodecVersion<1, 0, 0>,
 }
 
 impl Codec for FixedOffsetScaleCodec {
@@ -109,7 +112,7 @@ impl Codec for FixedOffsetScaleCodec {
 }
 
 impl StaticCodec for FixedOffsetScaleCodec {
-    const CODEC_ID: &'static str = "fixed-offset-scale";
+    const CODEC_ID: &'static str = "fixed-offset-scale.rs";
 
     type Config<'de> = Self;
 

codecs/fourier-network/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "numcodecs-fourier-network"
-version = "0.1.1"
+version = "0.2.0"
 edition = { workspace = true }
 authors = { workspace = true }
 repository = { workspace = true }

codecs/fourier-network/src/lib.rs

@@ -40,7 +40,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension, Ix1, O
 use num_traits::{ConstOne, ConstZero, Float as FloatTrait, FromPrimitive};
 use numcodecs::{
     AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
-    Codec, StaticCodec, StaticCodecConfig,
+    Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
 };
 use schemars::{json_schema, JsonSchema, Schema, SchemaGenerator};
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -59,6 +59,8 @@ mod modules;
 
 use modules::{Model, ModelConfig, ModelExtra, ModelRecord};
 
+type FourierNetworkCodecVersion = StaticCodecVersion<0, 1, 0>;
+
 #[derive(Clone, Serialize, Deserialize, JsonSchema)]
 #[serde(deny_unknown_fields)]
 /// Fourier network codec which trains and overfits a fourier feature neural
@@ -88,6 +90,9 @@ pub struct FourierNetworkCodec {
     pub mini_batch_size: Option<NonZeroUsize>,
     /// The seed for the random number generator used during encoding
     pub seed: u64,
+    /// The codec's encoding format version. Do not provide this parameter explicitly.
+    #[serde(default, rename = "_version")]
+    pub version: FourierNetworkCodecVersion,
 }
 
 // using this wrapper function makes an Option<T> required
@@ -176,7 +181,7 @@ impl Codec for FourierNetworkCodec {
 }
 
 impl StaticCodec for FourierNetworkCodec {
-    const CODEC_ID: &'static str = "fourier-network";
+    const CODEC_ID: &'static str = "fourier-network.rs";
 
     type Config<'de> = Self;
 
@@ -390,7 +395,7 @@ pub fn encode<T: FloatExt, S: Data<Elem = T>, D: Dimension, B: AutodiffBackend<F
     );
 
     let extra = ModelExtra {
-        model,
+        model: model.into_record(),
         b_t: Param::from_tensor(b_t).set_require_grad(false),
         mean: Param::from_tensor(Tensor::from_data(
             TensorData::new(vec![mean], vec![1]),
@@ -402,12 +407,11 @@ pub fn encode<T: FloatExt, S: Data<Elem = T>, D: Dimension, B: AutodiffBackend<F
             device,
         ))
         .set_require_grad(false),
+        version: StaticCodecVersion,
     };
 
     let recorder = BinBytesRecorder::<T::Precision>::new();
-    let encoded = recorder
-        .record(extra.into_record(), ())
-        .map_err(NeuralNetworkError)?;
+    let encoded = recorder.record(extra, ()).map_err(NeuralNetworkError)?;
 
     Ok(Array::from_vec(encoded))
 }
@@ -449,24 +453,14 @@ pub fn decode_into<T: FloatExt, S: Data<Elem = u8>, D: Dimension, B: Backend<Flo
     let encoded = encoded.into_owned().into_raw_vec_and_offset().0;
 
     let recorder = BinBytesRecorder::<T::Precision>::new();
-    let record = recorder.load(encoded, device).map_err(NeuralNetworkError)?;
-
-    let extra = ModelExtra::<B> {
-        model: ModelConfig::new(fourier_features, num_blocks).init(device),
-        b_t: Param::from_tensor(Tensor::zeros(
-            [decoded.ndim(), fourier_features.get()],
-            device,
-        ))
-        .set_require_grad(false),
-        mean: Param::from_tensor(Tensor::zeros([1], device)).set_require_grad(false),
-        stdv: Param::from_tensor(Tensor::ones([1], device)).set_require_grad(false),
-    }
-    .load_record(record);
-
-    let model = extra.model;
-    let b_t = extra.b_t.into_value();
-    let mean = extra.mean.into_value().into_scalar();
-    let stdv = extra.stdv.into_value().into_scalar();
+    let record: ModelExtra<B> = recorder.load(encoded, device).map_err(NeuralNetworkError)?;
+
+    let model = ModelConfig::new(fourier_features, num_blocks)
+        .init(device)
+        .load_record(record.model);
+    let b_t = record.b_t.into_value();
+    let mean = record.mean.into_value().into_scalar();
+    let stdv = record.stdv.into_value().into_scalar();
 
     let test_xs = flat_grid_like(&decoded, device);
     let test_xs = fourier_mapping(test_xs, b_t);

codecs/fourier-network/src/modules.rs

@@ -5,6 +5,7 @@ use burn::{
     module::{Module, Param},
     nn::{BatchNorm, BatchNormConfig, Gelu, Linear, LinearConfig},
     prelude::Backend,
+    record::{PrecisionSettings, Record},
     tensor::{Float, Tensor},
 };
 
@@ -93,10 +94,44 @@ impl ModelConfig {
     }
 }
 
-#[derive(Debug, Module)]
 pub struct ModelExtra<B: Backend> {
-    pub model: Model<B>,
+    pub model: <Model<B> as Module<B>>::Record,
     pub b_t: Param<Tensor<B, 2, Float>>,
     pub mean: Param<Tensor<B, 1, Float>>,
     pub stdv: Param<Tensor<B, 1, Float>>,
+    pub version: crate::FourierNetworkCodecVersion,
+}
+
+impl<B: Backend> Record<B> for ModelExtra<B> {
+    type Item<S: PrecisionSettings> = ModelExtraItem<B, S>;
+
+    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
+        ModelExtraItem {
+            model: self.model.into_item(),
+            b_t: self.b_t.into_item(),
+            mean: self.mean.into_item(),
+            stdv: self.stdv.into_item(),
+            version: self.version,
+        }
+    }
+
+    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
+        Self {
+            model: Record::<B>::from_item::<S>(item.model, device),
+            b_t: Record::<B>::from_item::<S>(item.b_t, device),
+            mean: Record::<B>::from_item::<S>(item.mean, device),
+            stdv: Record::<B>::from_item::<S>(item.stdv, device),
+            version: item.version,
+        }
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+#[serde(bound = "")]
+pub struct ModelExtraItem<B: Backend, S: PrecisionSettings> {
+    model: <<Model<B> as Module<B>>::Record as Record<B>>::Item<S>,
+    b_t: <Param<Tensor<B, 2, Float>> as Record<B>>::Item<S>,
+    mean: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
+    stdv: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
+    version: crate::FourierNetworkCodecVersion,
 }