Skip to content

Commit e098a6f

Browse files
juntyrjuntyr
authored
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-*
1 parent 7440b41 commit e098a6f

File tree

69 files changed

+949
-129
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

69 files changed

+949
-129
lines changed

‎Cargo.toml

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,8 @@ rust-version = "1.82"
4242

4343
[workspace.dependencies]
4444
# workspace-internal numcodecs crates
45-
numcodecs = { version = "0.2", path = "crates/numcodecs", default-features = false }
46-
numcodecs-python = { version = "0.5", path = "crates/numcodecs-python", default-features = false }
45+
numcodecs = { version = "0.2.1", path = "crates/numcodecs", default-features = false }
46+
numcodecs-python = { version = "0.6", path = "crates/numcodecs-python", default-features = false }
4747
numcodecs-wasm-builder = { version = "0.1", path = "crates/numcodecs-wasm-builder", default-features = false }
4848
numcodecs-wasm-logging = { version = "0.1", path = "crates/numcodecs-wasm-logging", default-features = false }
4949
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
5353
numcodecs-wasm = { version = "0.1", path = "py/numcodecs-wasm", default-features = false }
5454

5555
# workspace-internal codecs crates
56-
numcodecs-asinh = { version = "0.2", path = "codecs/asinh", default-features = false }
57-
numcodecs-bit-round = { version = "0.2", path = "codecs/bit-round", default-features = false }
58-
numcodecs-fixed-offset-scale = { version = "0.2", path = "codecs/fixed-offset-scale", default-features = false }
59-
numcodecs-fourier-network = { version = "0.1", path = "codecs/fourier-network", default-features = false }
60-
numcodecs-identity = { version = "0.2", path = "codecs/identity", default-features = false }
61-
numcodecs-jpeg2000 = { version = "0.1", path = "codecs/jpeg2000", default-features = false }
62-
numcodecs-linear-quantize = { version = "0.3", path = "codecs/linear-quantize", default-features = false }
63-
numcodecs-log = { version = "0.3", path = "codecs/log", default-features = false }
64-
numcodecs-pco = { version = "0.1", path = "codecs/pco", default-features = false }
65-
numcodecs-random-projection = { version = "0.2", path = "codecs/random-projection", default-features = false }
66-
numcodecs-reinterpret = { version = "0.2", path = "codecs/reinterpret", default-features = false }
67-
numcodecs-round = { version = "0.2", path = "codecs/round", default-features = false }
68-
numcodecs-swizzle-reshape = { version = "0.2", path = "codecs/swizzle-reshape", default-features = false }
69-
numcodecs-sz3 = { version = "0.5", path = "codecs/sz3", default-features = false }
70-
numcodecs-tthresh = { version = "0.1", path = "codecs/tthresh", default-features = false }
71-
numcodecs-uniform-noise = { version = "0.2", path = "codecs/uniform-noise", default-features = false }
72-
numcodecs-zfp = { version = "0.4", path = "codecs/zfp", default-features = false }
73-
numcodecs-zfp-classic = { version = "0.2", path = "codecs/zfp-classic", default-features = false }
74-
numcodecs-zlib = { version = "0.2", path = "codecs/zlib", default-features = false }
75-
numcodecs-zstd = { version = "0.2", path = "codecs/zstd", default-features = false }
56+
numcodecs-asinh = { version = "0.3", path = "codecs/asinh", default-features = false }
57+
numcodecs-bit-round = { version = "0.3", path = "codecs/bit-round", default-features = false }
58+
numcodecs-fixed-offset-scale = { version = "0.3", path = "codecs/fixed-offset-scale", default-features = false }
59+
numcodecs-fourier-network = { version = "0.2", path = "codecs/fourier-network", default-features = false }
60+
numcodecs-identity = { version = "0.3", path = "codecs/identity", default-features = false }
61+
numcodecs-jpeg2000 = { version = "0.2", path = "codecs/jpeg2000", default-features = false }
62+
numcodecs-linear-quantize = { version = "0.4", path = "codecs/linear-quantize", default-features = false }
63+
numcodecs-log = { version = "0.4", path = "codecs/log", default-features = false }
64+
numcodecs-pco = { version = "0.2", path = "codecs/pco", default-features = false }
65+
numcodecs-random-projection = { version = "0.3", path = "codecs/random-projection", default-features = false }
66+
numcodecs-reinterpret = { version = "0.3", path = "codecs/reinterpret", default-features = false }
67+
numcodecs-round = { version = "0.3", path = "codecs/round", default-features = false }
68+
numcodecs-swizzle-reshape = { version = "0.3", path = "codecs/swizzle-reshape", default-features = false }
69+
numcodecs-sz3 = { version = "0.6", path = "codecs/sz3", default-features = false }
70+
numcodecs-tthresh = { version = "0.2", path = "codecs/tthresh", default-features = false }
71+
numcodecs-uniform-noise = { version = "0.3", path = "codecs/uniform-noise", default-features = false }
72+
numcodecs-zfp = { version = "0.5", path = "codecs/zfp", default-features = false }
73+
numcodecs-zfp-classic = { version = "0.3", path = "codecs/zfp-classic", default-features = false }
74+
numcodecs-zlib = { version = "0.3", path = "codecs/zlib", default-features = false }
75+
numcodecs-zstd = { version = "0.3", path = "codecs/zstd", default-features = false }
7676

7777
# crates.io third-party dependencies
7878
anyhow = { version = "1.0.93", default-features = false }

‎codecs/asinh/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-asinh"
3-
version = "0.2.1"
3+
version = "0.3.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

‎codecs/asinh/src/lib.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension, Zip};
2121
use num_traits::{Float, Signed};
2222
use numcodecs::{
2323
AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
24-
Codec, StaticCodec, StaticCodecConfig,
24+
Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
2525
};
2626
use schemars::JsonSchema;
2727
use serde::{Deserialize, Serialize};
@@ -46,7 +46,10 @@ use thiserror::Error;
4646
pub struct AsinhCodec {
4747
/// The width of the close-to-zero input value range where the transform is
4848
/// nearly linear
49-
linear_width: f64,
49+
pub linear_width: f64,
50+
/// The codec's encoding format version. Do not provide this parameter explicitly.
51+
#[serde(default, rename = "_version")]
52+
pub version: StaticCodecVersion<1, 0, 0>,
5053
}
5154

5255
impl Codec for AsinhCodec {
@@ -99,7 +102,7 @@ impl Codec for AsinhCodec {
99102
}
100103

101104
impl StaticCodec for AsinhCodec {
102-
const CODEC_ID: &'static str = "asinh";
105+
const CODEC_ID: &'static str = "asinh.rs";
103106

104107
type Config<'de> = Self;
105108

‎codecs/bit-round/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-bit-round"
3-
version = "0.2.0"
3+
version = "0.3.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

‎codecs/bit-round/src/lib.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
use ndarray::{Array, ArrayBase, Data, Dimension};
2121
use numcodecs::{
2222
AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
23-
Codec, StaticCodec, StaticCodecConfig,
23+
Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
2424
};
2525
use schemars::JsonSchema;
2626
use serde::{Deserialize, Serialize};
@@ -45,6 +45,9 @@ pub struct BitRoundCodec {
4545
/// If keepbits is equal to the bitlength of the dtype's mantissa, no
4646
/// transformation is performed.
4747
pub keepbits: u8,
48+
/// The codec's encoding format version. Do not provide this parameter explicitly.
49+
#[serde(default, rename = "_version")]
50+
pub version: StaticCodecVersion<1, 0, 0>,
4851
}
4952

5053
impl Codec for BitRoundCodec {
@@ -80,7 +83,7 @@ impl Codec for BitRoundCodec {
8083
}
8184

8285
impl StaticCodec for BitRoundCodec {
83-
const CODEC_ID: &'static str = "bit-round";
86+
const CODEC_ID: &'static str = "bit-round.rs";
8487

8588
type Config<'de> = Self;
8689

‎codecs/fixed-offset-scale/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-fixed-offset-scale"
3-
version = "0.2.1"
3+
version = "0.3.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension};
2121
use num_traits::Float;
2222
use numcodecs::{
2323
AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
24-
Codec, StaticCodec, StaticCodecConfig,
24+
Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
2525
};
2626
use schemars::JsonSchema;
2727
use serde::{Deserialize, Serialize};
@@ -43,6 +43,9 @@ pub struct FixedOffsetScaleCodec {
4343
pub offset: f64,
4444
/// The scale of the data.
4545
pub scale: f64,
46+
/// The codec's encoding format version. Do not provide this parameter explicitly.
47+
#[serde(default, rename = "_version")]
48+
pub version: StaticCodecVersion<1, 0, 0>,
4649
}
4750

4851
impl Codec for FixedOffsetScaleCodec {
@@ -109,7 +112,7 @@ impl Codec for FixedOffsetScaleCodec {
109112
}
110113

111114
impl StaticCodec for FixedOffsetScaleCodec {
112-
const CODEC_ID: &'static str = "fixed-offset-scale";
115+
const CODEC_ID: &'static str = "fixed-offset-scale.rs";
113116

114117
type Config<'de> = Self;
115118

‎codecs/fourier-network/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-fourier-network"
3-
version = "0.1.1"
3+
version = "0.2.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

‎codecs/fourier-network/src/lib.rs

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension, Ix1, O
4040
use num_traits::{ConstOne, ConstZero, Float as FloatTrait, FromPrimitive};
4141
use numcodecs::{
4242
AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
43-
Codec, StaticCodec, StaticCodecConfig,
43+
Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
4444
};
4545
use schemars::{json_schema, JsonSchema, Schema, SchemaGenerator};
4646
use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -59,6 +59,8 @@ mod modules;
5959

6060
use modules::{Model, ModelConfig, ModelExtra, ModelRecord};
6161

62+
type FourierNetworkCodecVersion = StaticCodecVersion<0, 1, 0>;
63+
6264
#[derive(Clone, Serialize, Deserialize, JsonSchema)]
6365
#[serde(deny_unknown_fields)]
6466
/// Fourier network codec which trains and overfits a fourier feature neural
@@ -88,6 +90,9 @@ pub struct FourierNetworkCodec {
8890
pub mini_batch_size: Option<NonZeroUsize>,
8991
/// The seed for the random number generator used during encoding
9092
pub seed: u64,
93+
/// The codec's encoding format version. Do not provide this parameter explicitly.
94+
#[serde(default, rename = "_version")]
95+
pub version: FourierNetworkCodecVersion,
9196
}
9297

9398
// using this wrapper function makes an Option<T> required
@@ -176,7 +181,7 @@ impl Codec for FourierNetworkCodec {
176181
}
177182

178183
impl StaticCodec for FourierNetworkCodec {
179-
const CODEC_ID: &'static str = "fourier-network";
184+
const CODEC_ID: &'static str = "fourier-network.rs";
180185

181186
type Config<'de> = Self;
182187

@@ -390,7 +395,7 @@ pub fn encode<T: FloatExt, S: Data<Elem = T>, D: Dimension, B: AutodiffBackend<F
390395
);
391396

392397
let extra = ModelExtra {
393-
model,
398+
model: model.into_record(),
394399
b_t: Param::from_tensor(b_t).set_require_grad(false),
395400
mean: Param::from_tensor(Tensor::from_data(
396401
TensorData::new(vec![mean], vec![1]),
@@ -402,12 +407,11 @@ pub fn encode<T: FloatExt, S: Data<Elem = T>, D: Dimension, B: AutodiffBackend<F
402407
device,
403408
))
404409
.set_require_grad(false),
410+
version: StaticCodecVersion,
405411
};
406412

407413
let recorder = BinBytesRecorder::<T::Precision>::new();
408-
let encoded = recorder
409-
.record(extra.into_record(), ())
410-
.map_err(NeuralNetworkError)?;
414+
let encoded = recorder.record(extra, ()).map_err(NeuralNetworkError)?;
411415

412416
Ok(Array::from_vec(encoded))
413417
}
@@ -449,24 +453,14 @@ pub fn decode_into<T: FloatExt, S: Data<Elem = u8>, D: Dimension, B: Backend<Flo
449453
let encoded = encoded.into_owned().into_raw_vec_and_offset().0;
450454

451455
let recorder = BinBytesRecorder::<T::Precision>::new();
452-
let record = recorder.load(encoded, device).map_err(NeuralNetworkError)?;
453-
454-
let extra = ModelExtra::<B> {
455-
model: ModelConfig::new(fourier_features, num_blocks).init(device),
456-
b_t: Param::from_tensor(Tensor::zeros(
457-
[decoded.ndim(), fourier_features.get()],
458-
device,
459-
))
460-
.set_require_grad(false),
461-
mean: Param::from_tensor(Tensor::zeros([1], device)).set_require_grad(false),
462-
stdv: Param::from_tensor(Tensor::ones([1], device)).set_require_grad(false),
463-
}
464-
.load_record(record);
465-
466-
let model = extra.model;
467-
let b_t = extra.b_t.into_value();
468-
let mean = extra.mean.into_value().into_scalar();
469-
let stdv = extra.stdv.into_value().into_scalar();
456+
let record: ModelExtra<B> = recorder.load(encoded, device).map_err(NeuralNetworkError)?;
457+
458+
let model = ModelConfig::new(fourier_features, num_blocks)
459+
.init(device)
460+
.load_record(record.model);
461+
let b_t = record.b_t.into_value();
462+
let mean = record.mean.into_value().into_scalar();
463+
let stdv = record.stdv.into_value().into_scalar();
470464

471465
let test_xs = flat_grid_like(&decoded, device);
472466
let test_xs = fourier_mapping(test_xs, b_t);

‎codecs/fourier-network/src/modules.rs

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use burn::{
55
module::{Module, Param},
66
nn::{BatchNorm, BatchNormConfig, Gelu, Linear, LinearConfig},
77
prelude::Backend,
8+
record::{PrecisionSettings, Record},
89
tensor::{Float, Tensor},
910
};
1011

@@ -93,10 +94,44 @@ impl ModelConfig {
9394
}
9495
}
9596

96-
#[derive(Debug, Module)]
9797
pub struct ModelExtra<B: Backend> {
98-
pub model: Model<B>,
98+
pub model: <Model<B> as Module<B>>::Record,
9999
pub b_t: Param<Tensor<B, 2, Float>>,
100100
pub mean: Param<Tensor<B, 1, Float>>,
101101
pub stdv: Param<Tensor<B, 1, Float>>,
102+
pub version: crate::FourierNetworkCodecVersion,
103+
}
104+
105+
impl<B: Backend> Record<B> for ModelExtra<B> {
106+
type Item<S: PrecisionSettings> = ModelExtraItem<B, S>;
107+
108+
fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
109+
ModelExtraItem {
110+
model: self.model.into_item(),
111+
b_t: self.b_t.into_item(),
112+
mean: self.mean.into_item(),
113+
stdv: self.stdv.into_item(),
114+
version: self.version,
115+
}
116+
}
117+
118+
fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
119+
Self {
120+
model: Record::<B>::from_item::<S>(item.model, device),
121+
b_t: Record::<B>::from_item::<S>(item.b_t, device),
122+
mean: Record::<B>::from_item::<S>(item.mean, device),
123+
stdv: Record::<B>::from_item::<S>(item.stdv, device),
124+
version: item.version,
125+
}
126+
}
127+
}
128+
129+
#[derive(serde::Serialize, serde::Deserialize)]
130+
#[serde(bound = "")]
131+
pub struct ModelExtraItem<B: Backend, S: PrecisionSettings> {
132+
model: <<Model<B> as Module<B>>::Record as Record<B>>::Item<S>,
133+
b_t: <Param<Tensor<B, 2, Float>> as Record<B>>::Item<S>,
134+
mean: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
135+
stdv: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
136+
version: crate::FourierNetworkCodecVersion,
102137
}

‎codecs/fourier-network/tests/schema.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,12 @@
4444
"format": "uint64",
4545
"minimum": 0,
4646
"description": "The seed for the random number generator used during encoding"
47+
},
48+
"_version": {
49+
"type": "string",
50+
"pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
51+
"description": "The codec's encoding format version. Do not provide this parameter explicitly.",
52+
"default": "0.1.0"
4753
}
4854
},
4955
"required": [

‎codecs/identity/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-identity"
3-
version = "0.2.0"
3+
version = "0.3.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

‎codecs/identity/src/lib.rs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
2020
use numcodecs::{
2121
AnyArray, AnyArrayAssignError, AnyArrayView, AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
22-
StaticCodecConfig,
22+
StaticCodecConfig, StaticCodecVersion,
2323
};
2424
use schemars::JsonSchema;
2525
use serde::{Deserialize, Serialize};
@@ -30,7 +30,9 @@ use thiserror::Error;
3030
/// Identity codec which applies the identity function, i.e. passes through the
3131
/// input unchanged during encoding and decoding.
3232
pub struct IdentityCodec {
33-
// empty
33+
/// The codec's encoding format version. Do not provide this parameter explicitly.
34+
#[serde(default, rename = "_version")]
35+
pub version: StaticCodecVersion<1, 0, 0>,
3436
}
3537

3638
impl Codec for IdentityCodec {
@@ -54,7 +56,7 @@ impl Codec for IdentityCodec {
5456
}
5557

5658
impl StaticCodec for IdentityCodec {
57-
const CODEC_ID: &'static str = "identity";
59+
const CODEC_ID: &'static str = "identity.rs";
5860

5961
type Config<'de> = Self;
6062

‎codecs/jpeg2000/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "numcodecs-jpeg2000"
3-
version = "0.1.1"
3+
version = "0.2.0"
44
edition = { workspace = true }
55
authors = { workspace = true }
66
repository = { workspace = true }

0 commit comments

Comments
 (0)