j2k_ml/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Thin Burn tensor adapter for the `j2k` owned batch codec.
4//!
5//! The codec crates own parsing, grouping, decoding, and accelerator execution.
6//! This crate materializes CPU groups or stages completed CUDA and Metal codec
7//! output through host memory before an ordinary Burn upload. Casting and
8//! normalization stay in ordinary Burn tensor operations after decode.
9
10#![deny(missing_docs)]
11
12use burn_core::tensor::{backend::Backend, Int, Tensor};
13use j2k::{BatchGroupInfo, IndexedBatchError, J2kDecodeWarning, Rect};
14
15#[cfg(any(
16    feature = "cpu",
17    feature = "cuda",
18    all(feature = "metal", target_os = "macos")
19))]
20mod batch_contract;
21#[cfg(test)]
22mod completion;
23mod error;
24#[cfg(any(feature = "cuda", all(feature = "metal", target_os = "macos")))]
25mod staging;
26
27#[cfg(feature = "cpu")]
28pub mod cpu;
29#[cfg(feature = "cuda")]
30pub mod cuda;
31#[cfg(feature = "metal")]
32pub mod metal;
33
34#[cfg(feature = "cpu")]
35pub use cpu::CpuBurnDecoder;
36#[cfg(feature = "cuda")]
37pub use cuda::{
38    CudaBurnDecoder, CudaUploadBurnDecoder, SubmittedCudaBurnBatch, SubmittedCudaUploadBurnBatch,
39};
40pub use error::BurnDecodeError;
41#[cfg(feature = "metal")]
42pub use metal::{
43    MetalBurnDecoder, MetalUploadBurnDecoder, SubmittedMetalBurnBatch,
44    SubmittedMetalUploadBurnBatch,
45};
46
47/// Ordinary rank-4 Burn integer tensor tagged with its exact codec sample type.
48#[derive(Debug)]
49pub enum BurnBatchTensor<B: Backend> {
50    /// Unsigned samples with precision at most eight bits.
51    U8(Tensor<B, 4, Int>),
52    /// Unsigned samples with precision from nine through sixteen bits.
53    U16(Tensor<B, 4, Int>),
54    /// Signed samples with precision at most sixteen bits.
55    I16(Tensor<B, 4, Int>),
56}
57
58impl<B: Backend> BurnBatchTensor<B> {
59    /// Borrow the ordinary Burn integer tensor regardless of its codec dtype tag.
60    #[must_use]
61    pub const fn tensor(&self) -> &Tensor<B, 4, Int> {
62        match self {
63            Self::U8(tensor) | Self::U16(tensor) | Self::I16(tensor) => tensor,
64        }
65    }
66
67    /// Consume the codec dtype tag and return the ordinary Burn integer tensor.
68    #[must_use]
69    pub fn into_tensor(self) -> Tensor<B, 4, Int> {
70        match self {
71            Self::U8(tensor) | Self::U16(tensor) | Self::I16(tensor) => tensor,
72        }
73    }
74}
75
76/// One homogeneous decoded tensor group and its codec metadata.
77#[derive(Debug)]
78pub struct BurnBatchGroup<B: Backend> {
79    /// Decoded rank-4 integer tensor.
80    pub tensor: BurnBatchTensor<B>,
81    /// Exact codec and output metadata shared by the group.
82    pub info: BatchGroupInfo,
83    /// Original caller indices in tensor batch order.
84    pub source_indices: Vec<usize>,
85    /// Actual decoded source rectangle for each tensor item.
86    pub decoded_rects: Vec<Rect>,
87    /// Non-fatal codec warnings for each tensor item.
88    pub warnings: Vec<Vec<J2kDecodeWarning>>,
89}
90
91/// Failure while submitting or completing one homogeneous Burn tensor group.
92///
93/// No partially written tensor from the affected group is exposed. Other
94/// homogeneous groups may still succeed when the retained codec and framework
95/// sessions remain usable.
96#[derive(Debug, thiserror::Error)]
97#[error("Burn batch group containing source indices {source_indices:?} failed: {source}")]
98pub struct BurnBatchGroupError {
99    source_indices: Vec<usize>,
100    #[source]
101    source: BurnDecodeError,
102}
103
104impl BurnBatchGroupError {
105    #[cfg(any(feature = "cuda", all(feature = "metal", target_os = "macos"), test))]
106    pub(crate) fn new(source_indices: Vec<usize>, source: BurnDecodeError) -> Self {
107        Self {
108            source_indices,
109            source,
110        }
111    }
112
113    /// Original input indices whose dense tensor group was discarded.
114    #[must_use]
115    pub fn source_indices(&self) -> &[usize] {
116        &self.source_indices
117    }
118
119    /// Structured codec, framework, or interop failure for this group.
120    #[must_use]
121    pub const fn source(&self) -> &BurnDecodeError {
122        &self.source
123    }
124
125    /// Consume the group failure into affected indices and its source.
126    #[must_use]
127    pub fn into_parts(self) -> (Vec<usize>, BurnDecodeError) {
128        (self.source_indices, self.source)
129    }
130}
131
132/// Successful tensor groups plus indexed preparation and homogeneous execution failures.
133#[derive(Debug)]
134pub struct BurnBatchDecode<B: Backend> {
135    /// Successfully decoded homogeneous tensor groups.
136    pub groups: Vec<BurnBatchGroup<B>>,
137    /// Structured preparation or decode failures in original input order.
138    pub errors: Vec<IndexedBatchError>,
139    /// Homogeneous groups discarded after adapter submission or completion failed.
140    pub group_errors: Vec<BurnBatchGroupError>,
141}
j2k_ml/lib.rs

j2k_ml/
lib.rs
