Skip to main content

konjure_sdk/data/
inspect.rs

1//! Bounded, serializable summaries of immutable tensors.
2
3use alloc::{string::ToString, vec, vec::Vec};
4use core::cmp::Ordering;
5use serde::{Deserialize, Serialize};
6
7use super::{DType, DataError, Scalar, Tensor, wire};
8
9/// The largest number of logical tensor elements one inspection will read.
10pub const MAX_INSPECTION_SAMPLES: usize = 4_096;
11/// The largest exact prefix included in a tensor inspection.
12pub const MAX_INSPECTION_PREVIEW: usize = 32;
13/// The number of buckets used for nonconstant approximate histograms.
14pub const HISTOGRAM_BUCKETS: usize = 16;
15
16/// Bounds for one [`TensorInspection`].
17///
18/// `max_samples` limits all logical element reads, including the preview. It
19/// is capped at [`MAX_INSPECTION_SAMPLES`] so a host cannot accidentally turn
20/// an inspector click into an unbounded traversal.
21#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
22#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
23#[serde(deny_unknown_fields)]
24pub struct TensorInspectionOptions {
25    /// Maximum number of logical elements to inspect; zero produces metadata only.
26    pub max_samples: usize,
27}
28
29impl Default for TensorInspectionOptions {
30    fn default() -> Self {
31        Self {
32            max_samples: MAX_INSPECTION_SAMPLES,
33        }
34    }
35}
36
37/// One compact, approximate histogram interval.
38#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
39#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
40#[serde(deny_unknown_fields)]
41pub struct HistogramBucket {
42    /// Inclusive lower bound in the approximate binary64 analysis domain.
43    pub start: f64,
44    /// Inclusive upper bound for the final bucket and exclusive otherwise.
45    pub end: f64,
46    /// Number of inspected values in this interval.
47    pub count: usize,
48}
49
50/// Numerically stable statistics calculated from bounded logical samples.
51#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
52#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
53#[serde(deny_unknown_fields)]
54pub struct ApproximateTensorStatistics {
55    /// Arithmetic mean computed by Welford's online method.
56    pub mean: f64,
57    /// Population standard deviation computed by Welford's online method.
58    pub standard_deviation: f64,
59    /// Histogram over the values used for the mean and standard deviation.
60    pub histogram: Vec<HistogramBucket>,
61}
62
63/// Why binary64 approximate statistics could not be represented honestly.
64#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
65#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
66#[serde(rename_all = "snake_case")]
67pub enum ApproximationUnavailable {
68    /// No samples were requested or the tensor has no logical elements.
69    NoSamples,
70    /// At least one exact scalar lies outside the finite binary64 analysis domain.
71    ValueOutOfRange,
72    /// Finite binary64 values had a range or variance too wide for binary64 arithmetic.
73    ArithmeticOverflow,
74}
75
76/// Approximate statistics, or an explicit reason they are unavailable.
77#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
78#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
79#[serde(tag = "kind", rename_all = "snake_case")]
80pub enum TensorStatistics {
81    /// Binary64 statistics over the reported sample set.
82    Available(ApproximateTensorStatistics),
83    /// Statistics were omitted rather than rounded into a misleading result.
84    Unavailable {
85        /// Specific representability or sampling condition that prevented analysis.
86        reason: ApproximationUnavailable,
87    },
88}
89
90/// A bounded inspection of a tensor's logical values and retained storage.
91///
92/// `minimum`, `maximum`, and `preview` retain their exact declared scalar
93/// representation. Extrema cover the whole tensor when
94/// `sampled_values == logical_elements`; otherwise they describe the deterministic
95/// sample set. Preview values always show the exact leading logical elements. Approximate statistics
96/// use the same sample set and state explicitly when binary64 cannot represent it.
97#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
98#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
99#[serde(deny_unknown_fields)]
100pub struct TensorInspection {
101    /// Declared element representation.
102    pub dtype: DType,
103    /// Logical dimensions in axis order.
104    pub shape: Vec<usize>,
105    /// Logical element strides in axis order; negative strides identify reverse views.
106    pub strides: Vec<isize>,
107    /// Number of logical elements in the tensor view.
108    pub logical_elements: usize,
109    /// Bytes occupied by the logical elements if materialized contiguously.
110    pub logical_bytes: usize,
111    /// Bytes retained by the shared packed backing allocation.
112    pub retained_bytes: usize,
113    /// Number of logical values actually read for this inspection.
114    pub sampled_values: usize,
115    /// Whether every logical value was read, making extrema exact for the tensor.
116    pub extrema_complete: bool,
117    /// First logical values, up to [`MAX_INSPECTION_PREVIEW`] and the sample limit.
118    pub preview: Vec<wire::Scalar>,
119    /// Exact scalar minimum over all logical values or the deterministic sample set.
120    pub minimum: Option<wire::Scalar>,
121    /// Exact scalar maximum over all logical values or the deterministic sample set.
122    pub maximum: Option<wire::Scalar>,
123    /// Finite binary64 analysis of the deterministic sample set.
124    pub statistics: TensorStatistics,
125}
126
127impl Tensor {
128    /// Inspects this tensor without cloning or materializing its backing storage.
129    ///
130    /// Views are visited in logical row-major order through their shape and
131    /// strides. The inspection reads no more than `options.max_samples`, capped
132    /// by [`MAX_INSPECTION_SAMPLES`].
133    pub fn inspect(&self, options: TensorInspectionOptions) -> Result<TensorInspection, DataError> {
134        inspect_tensor(self, options)
135    }
136}
137
138/// Inspects a tensor without cloning or materializing its backing storage.
139///
140/// # Errors
141///
142/// Propagates a layout bounds error if a tensor invariant has been violated.
143pub fn inspect_tensor(
144    tensor: &Tensor,
145    options: TensorInspectionOptions,
146) -> Result<TensorInspection, DataError> {
147    let logical_elements = tensor.len();
148    let sample_count = logical_elements.min(options.max_samples.min(MAX_INSPECTION_SAMPLES));
149    let positions = sample_positions(logical_elements, sample_count);
150    let mut preview = Vec::with_capacity(positions.len().min(MAX_INSPECTION_PREVIEW));
151    let mut minimum = None;
152    let mut maximum = None;
153    let mut analysis = Vec::with_capacity(positions.len());
154
155    for (sample_index, position) in positions.into_iter().enumerate() {
156        let value = logical_value(tensor, position)?;
157        if sample_index < MAX_INSPECTION_PREVIEW {
158            preview.push(wire_scalar(value));
159        }
160        if minimum.is_none_or(|current| value.compare(current) == Ok(Ordering::Less)) {
161            minimum = Some(value);
162        }
163        if maximum.is_none_or(|current| value.compare(current) == Ok(Ordering::Greater)) {
164            maximum = Some(value);
165        }
166        analysis.push(value);
167    }
168
169    let sampled_values = analysis.len();
170    Ok(TensorInspection {
171        dtype: tensor.dtype(),
172        shape: tensor.shape().to_vec(),
173        strides: tensor.strides().to_vec(),
174        logical_elements,
175        logical_bytes: tensor.byte_len(),
176        retained_bytes: tensor.storage_byte_len(),
177        sampled_values,
178        extrema_complete: sampled_values == logical_elements,
179        preview,
180        minimum: minimum.map(wire_scalar),
181        maximum: maximum.map(wire_scalar),
182        statistics: statistics(&analysis),
183    })
184}
185
186fn sample_positions(len: usize, count: usize) -> Vec<usize> {
187    let preview_count = count.min(MAX_INSPECTION_PREVIEW);
188    let mut positions = Vec::with_capacity(count);
189    positions.extend(0..preview_count);
190    let remaining = count - preview_count;
191    if remaining == 0 {
192        return positions;
193    }
194    let tail_len = len - preview_count;
195    for index in 0..remaining {
196        positions.push(preview_count + evenly_spaced(index, remaining, tail_len));
197    }
198    positions
199}
200
201fn evenly_spaced(index: usize, count: usize, len: usize) -> usize {
202    if count <= 1 {
203        return len.saturating_sub(1);
204    }
205    let denominator = count - 1;
206    let numerator = len - 1;
207    // `len` is bounded by Tensor::MAX_ELEMENTS and `count` by 4096, but split
208    // the product anyway so this helper remains overflow-free if those bounds grow.
209    let quotient = numerator / denominator;
210    let remainder = numerator % denominator;
211    quotient * index + remainder * index / denominator
212}
213
214fn logical_value(tensor: &Tensor, mut flat: usize) -> Result<Scalar, DataError> {
215    let mut indices = Vec::with_capacity(tensor.shape().len());
216    for &dimension in tensor.shape().iter().rev() {
217        indices.push(isize::try_from(flat % dimension).map_err(|_| DataError::Bounds)?);
218        flat /= dimension;
219    }
220    indices.reverse();
221    tensor.get(&indices)
222}
223
224fn wire_scalar(value: Scalar) -> wire::Scalar {
225    wire::Scalar {
226        dtype: value.dtype(),
227        value: value.to_string(),
228    }
229}
230
231fn statistics(values: &[Scalar]) -> TensorStatistics {
232    if values.is_empty() {
233        return TensorStatistics::Unavailable {
234            reason: ApproximationUnavailable::NoSamples,
235        };
236    }
237    let mut converted = Vec::with_capacity(values.len());
238    for &value in values {
239        let Ok(approximate) = value.to_string().parse::<f64>() else {
240            return TensorStatistics::Unavailable {
241                reason: ApproximationUnavailable::ValueOutOfRange,
242            };
243        };
244        if !approximate.is_finite() {
245            return TensorStatistics::Unavailable {
246                reason: ApproximationUnavailable::ValueOutOfRange,
247            };
248        }
249        if approximate == 0.0 && value.compare(Scalar::zero(value.dtype())) != Ok(Ordering::Equal) {
250            // Parsing a finite binary128 subnormal into binary64 can underflow
251            // to zero. Preserve the exact scalar in the inspection and decline
252            // binary64 statistics rather than treating it as a true zero.
253            return TensorStatistics::Unavailable {
254                reason: ApproximationUnavailable::ValueOutOfRange,
255            };
256        }
257        converted.push(approximate);
258    }
259
260    let scale = converted.iter().copied().map(f64::abs).fold(0.0, f64::max);
261    if scale == 0.0 {
262        return TensorStatistics::Available(ApproximateTensorStatistics {
263            mean: 0.0,
264            standard_deviation: 0.0,
265            histogram: vec![HistogramBucket {
266                start: 0.0,
267                end: 0.0,
268                count: converted.len(),
269            }],
270        });
271    }
272
273    // Scale first so Welford's second moment neither overflows for values near
274    // f64::MAX nor underflows for small, representable finite values.
275    let mut mean: f64 = 0.0;
276    let mut sum_of_squares: f64 = 0.0;
277    for (index, value) in converted
278        .iter()
279        .copied()
280        .map(|value| value / scale)
281        .enumerate()
282    {
283        let count = (index + 1) as f64;
284        let delta = value - mean;
285        mean += delta / count;
286        sum_of_squares += delta * (value - mean);
287        if !mean.is_finite() || !sum_of_squares.is_finite() {
288            return TensorStatistics::Unavailable {
289                reason: ApproximationUnavailable::ArithmeticOverflow,
290            };
291        }
292    }
293    let variance = sum_of_squares / converted.len() as f64;
294    if !variance.is_finite() || variance < 0.0 {
295        return TensorStatistics::Unavailable {
296            reason: ApproximationUnavailable::ArithmeticOverflow,
297        };
298    }
299    let mean = mean * scale;
300    let standard_deviation = libm::sqrt(variance) * scale;
301    if !mean.is_finite() || !standard_deviation.is_finite() {
302        return TensorStatistics::Unavailable {
303            reason: ApproximationUnavailable::ArithmeticOverflow,
304        };
305    }
306    let histogram = histogram(&converted, scale);
307    TensorStatistics::Available(ApproximateTensorStatistics {
308        mean,
309        standard_deviation,
310        histogram,
311    })
312}
313
314fn histogram(values: &[f64], scale: f64) -> Vec<HistogramBucket> {
315    let Some((&first, rest)) = values.split_first() else {
316        return Vec::new();
317    };
318    let (minimum, maximum) = rest
319        .iter()
320        .copied()
321        .fold((first, first), |(minimum, maximum), value| {
322            (minimum.min(value), maximum.max(value))
323        });
324    if minimum == maximum {
325        return vec![HistogramBucket {
326            start: minimum,
327            end: maximum,
328            count: values.len(),
329        }];
330    }
331    let normalized_minimum = minimum / scale;
332    let normalized_maximum = maximum / scale;
333    let width = (normalized_maximum - normalized_minimum) / HISTOGRAM_BUCKETS as f64;
334    if !width.is_finite() || width <= 0.0 {
335        return vec![HistogramBucket {
336            start: minimum,
337            end: maximum,
338            count: values.len(),
339        }];
340    }
341    let mut counts = vec![0usize; HISTOGRAM_BUCKETS];
342    for value in values {
343        let index =
344            (((*value / scale - normalized_minimum) / width) as usize).min(HISTOGRAM_BUCKETS - 1);
345        counts[index] += 1;
346    }
347    let mut boundaries = Vec::with_capacity(HISTOGRAM_BUCKETS + 1);
348    for index in 0..=HISTOGRAM_BUCKETS {
349        boundaries.push(if index == 0 {
350            minimum
351        } else if index == HISTOGRAM_BUCKETS {
352            maximum
353        } else {
354            (normalized_minimum + width * index as f64) * scale
355        });
356    }
357    if boundaries
358        .windows(2)
359        .any(|pair| !pair[0].is_finite() || !pair[1].is_finite() || pair[0] >= pair[1])
360    {
361        // The endpoints cannot support distinct representable intervals (most
362        // commonly a subnormal range). A single exact range is more honest than
363        // emitting repeated or unordered approximate bucket boundaries.
364        return vec![HistogramBucket {
365            start: minimum,
366            end: maximum,
367            count: values.len(),
368        }];
369    }
370    counts
371        .into_iter()
372        .enumerate()
373        .map(|(index, count)| HistogramBucket {
374            start: boundaries[index],
375            end: boundaries[index + 1],
376            count,
377        })
378        .collect()
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384    use crate::data::AxisIndex;
385
386    fn tensor(dtype: DType, values: &[&str]) -> Tensor {
387        Tensor::from_scalars(
388            dtype,
389            vec![values.len()],
390            values
391                .iter()
392                .map(|value| Scalar::parse(dtype, value).unwrap())
393                .collect(),
394        )
395        .unwrap()
396    }
397
398    #[test]
399    fn empty_tensor_reports_metadata_without_statistics() {
400        let inspection = Tensor::zeros(DType::F32, vec![0, 4])
401            .unwrap()
402            .inspect(TensorInspectionOptions::default())
403            .unwrap();
404        assert_eq!(inspection.logical_elements, 0);
405        assert!(inspection.preview.is_empty());
406        assert_eq!(inspection.minimum, None);
407        assert!(matches!(
408            inspection.statistics,
409            TensorStatistics::Unavailable {
410                reason: ApproximationUnavailable::NoSamples
411            }
412        ));
413    }
414
415    #[test]
416    fn constant_tensor_has_one_complete_histogram_bucket() {
417        let inspection = tensor(DType::I16, &["7", "7", "7"])
418            .inspect(TensorInspectionOptions::default())
419            .unwrap();
420        let TensorStatistics::Available(statistics) = inspection.statistics else {
421            panic!("constant values should be representable")
422        };
423        assert_eq!(statistics.mean, 7.0);
424        assert_eq!(statistics.standard_deviation, 0.0);
425        assert_eq!(statistics.histogram.len(), 1);
426        assert_eq!(statistics.histogram[0].count, 3);
427    }
428
429    #[test]
430    fn wide_integer_extrema_remain_exact_while_statistics_are_approximate() {
431        let inspection = tensor(
432            DType::U128,
433            &["0", "340282366920938463463374607431768211455"],
434        )
435        .inspect(TensorInspectionOptions::default())
436        .unwrap();
437        assert_eq!(inspection.minimum.unwrap().value, "0");
438        assert_eq!(
439            inspection.maximum.unwrap().value,
440            "340282366920938463463374607431768211455"
441        );
442        assert!(matches!(
443            inspection.statistics,
444            TensorStatistics::Available(_)
445        ));
446    }
447
448    #[test]
449    fn reverse_view_is_read_in_logical_order() {
450        let view = tensor(DType::I8, &["1", "2", "3", "4"])
451            .slice(&[AxisIndex::Slice {
452                start: None,
453                stop: None,
454                step: -1,
455            }])
456            .unwrap();
457        let inspection = view.inspect(TensorInspectionOptions::default()).unwrap();
458        assert_eq!(inspection.strides, vec![-1]);
459        assert_eq!(
460            inspection
461                .preview
462                .iter()
463                .map(|value| value.value.as_str())
464                .collect::<Vec<_>>(),
465            ["4", "3", "2", "1"]
466        );
467        assert_eq!(inspection.minimum.unwrap().value, "1");
468        assert_eq!(inspection.maximum.unwrap().value, "4");
469    }
470
471    #[test]
472    fn bounded_sampling_includes_prefix_and_last_value() {
473        let values = (0..100).map(|value| value.to_string()).collect::<Vec<_>>();
474        let scalars = values
475            .iter()
476            .map(|value| Scalar::parse(DType::I16, value).unwrap())
477            .collect();
478        let inspection = Tensor::from_scalars(DType::I16, vec![100], scalars)
479            .unwrap()
480            .inspect(TensorInspectionOptions { max_samples: 33 })
481            .unwrap();
482        assert_eq!(inspection.sampled_values, 33);
483        assert!(!inspection.extrema_complete);
484        assert_eq!(inspection.preview.len(), 32);
485        assert_eq!(inspection.preview[0].value, "0");
486        assert_eq!(inspection.maximum.unwrap().value, "99");
487    }
488
489    #[test]
490    fn f128_outside_binary64_reports_unavailable_statistics() {
491        let inspection = tensor(DType::F128, &["1e400", "2e400"])
492            .inspect(TensorInspectionOptions::default())
493            .unwrap();
494        assert!(inspection.minimum.unwrap().value.ends_with("E+400"));
495        assert!(matches!(
496            inspection.statistics,
497            TensorStatistics::Unavailable {
498                reason: ApproximationUnavailable::ValueOutOfRange
499            }
500        ));
501    }
502
503    #[test]
504    fn negative_and_subnormal_values_have_a_finite_histogram() {
505        let inspection = tensor(DType::F64, &["-2", "-5e-324", "5e-324", "2"])
506            .inspect(TensorInspectionOptions::default())
507            .unwrap();
508        let TensorStatistics::Available(statistics) = inspection.statistics else {
509            panic!("finite f64 values should be analysable")
510        };
511        assert_eq!(
512            statistics
513                .histogram
514                .iter()
515                .map(|bucket| bucket.count)
516                .sum::<usize>(),
517            4
518        );
519    }
520
521    #[test]
522    fn f128_underflow_to_binary64_zero_is_unavailable() {
523        let inspection = tensor(DType::F128, &["1e-400"])
524            .inspect(TensorInspectionOptions::default())
525            .unwrap();
526        assert!(matches!(
527            inspection.statistics,
528            TensorStatistics::Unavailable {
529                reason: ApproximationUnavailable::ValueOutOfRange
530            }
531        ));
532    }
533
534    #[test]
535    fn scaled_statistics_preserve_small_standard_deviation() {
536        let inspection = tensor(DType::F64, &["-1e-300", "1e-300"])
537            .inspect(TensorInspectionOptions::default())
538            .unwrap();
539        let TensorStatistics::Available(statistics) = inspection.statistics else {
540            panic!("representable values should be analysable")
541        };
542        assert!(statistics.mean.abs() < 1e-315);
543        assert!((statistics.standard_deviation - 1e-300).abs() < 1e-312);
544        assert_histogram_is_ordered(&statistics.histogram, 2);
545    }
546
547    #[test]
548    fn scaled_statistics_preserve_wide_finite_range() {
549        let inspection = tensor(DType::F64, &["-1e308", "1e308"])
550            .inspect(TensorInspectionOptions::default())
551            .unwrap();
552        let TensorStatistics::Available(statistics) = inspection.statistics else {
553            panic!("representable values should be analysable")
554        };
555        assert!(statistics.mean.abs() < 1e293);
556        assert!((statistics.standard_deviation - 1e308).abs() < 1e295);
557        assert_histogram_is_ordered(&statistics.histogram, 2);
558    }
559
560    fn assert_histogram_is_ordered(histogram: &[HistogramBucket], count: usize) {
561        assert_eq!(
562            histogram.iter().map(|bucket| bucket.count).sum::<usize>(),
563            count
564        );
565        assert!(histogram.iter().all(|bucket| {
566            bucket.start.is_finite() && bucket.end.is_finite() && bucket.start <= bucket.end
567        }));
568        assert!(
569            histogram
570                .windows(2)
571                .all(|pair| pair[0].end <= pair[1].start)
572        );
573    }
574}