1use alloc::{string::ToString, vec, vec::Vec};
4use core::cmp::Ordering;
5use serde::{Deserialize, Serialize};
6
7use super::{DType, DataError, Scalar, Tensor, wire};
8
9pub const MAX_INSPECTION_SAMPLES: usize = 4_096;
11pub const MAX_INSPECTION_PREVIEW: usize = 32;
13pub const HISTOGRAM_BUCKETS: usize = 16;
15
16#[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 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#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
39#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
40#[serde(deny_unknown_fields)]
41pub struct HistogramBucket {
42 pub start: f64,
44 pub end: f64,
46 pub count: usize,
48}
49
50#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
52#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
53#[serde(deny_unknown_fields)]
54pub struct ApproximateTensorStatistics {
55 pub mean: f64,
57 pub standard_deviation: f64,
59 pub histogram: Vec<HistogramBucket>,
61}
62
63#[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 NoSamples,
70 ValueOutOfRange,
72 ArithmeticOverflow,
74}
75
76#[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 Available(ApproximateTensorStatistics),
83 Unavailable {
85 reason: ApproximationUnavailable,
87 },
88}
89
90#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
98#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
99#[serde(deny_unknown_fields)]
100pub struct TensorInspection {
101 pub dtype: DType,
103 pub shape: Vec<usize>,
105 pub strides: Vec<isize>,
107 pub logical_elements: usize,
109 pub logical_bytes: usize,
111 pub retained_bytes: usize,
113 pub sampled_values: usize,
115 pub extrema_complete: bool,
117 pub preview: Vec<wire::Scalar>,
119 pub minimum: Option<wire::Scalar>,
121 pub maximum: Option<wire::Scalar>,
123 pub statistics: TensorStatistics,
125}
126
127impl Tensor {
128 pub fn inspect(&self, options: TensorInspectionOptions) -> Result<TensorInspection, DataError> {
134 inspect_tensor(self, options)
135 }
136}
137
138pub 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 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 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 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 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}