1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
//! You can register hooks on `ag::Tensor` objects for debugging.
//! ```
//! use scirs2_autograd as ag;
//! use ag::tensor_ops as T;
//!
//! ag::run(|ctx| {
//! let a: ag::Tensor<f32> = T::zeros(&[4, 2], ctx).show();
//! let b: ag::Tensor<f32> = T::ones(&[2, 3], ctx).showshape();
//! let c = T::matmul(a, b).show_prefixed("MatMul:");
//!
//! c.eval( ctx);
//! // [[0.0, 0.0],
//! // [0.0, 0.0],
//! // [0.0, 0.0],
//! // [0.0, 0.0]] shape=[4, 2], strides=[2, 1], layout=C (0x1)
//! //
//! // [2, 3]
//! //
//! // MatMul:
//! // [[0.0, 0.0, 0.0],
//! // [0.0, 0.0, 0.0],
//! // [0.0, 0.0, 0.0],
//! // [0.0, 0.0, 0.0]] shape=[4, 3], strides=[3, 1], layout=C (0x1), dynamic ndim=2
//!
//! // raw hook
//! let a: ag::Tensor<f32> = T::zeros(&[4, 2], ctx)
//! .raw_hook(|x| println!("{:?}", x.shape()));
//! });
//!
//! ```
use super::*;
use crate::ndarray_ext::NdArrayView;
use std::marker::PhantomData;
/// Trait for hooks
///
/// hooks can be set using [crate::tensor::Tensor::register_hook()] method.
pub trait Hook<T: Float> {
fn call(&self, arr: &crate::ndarray::ArrayViewD<T>);
}
pub struct Print(pub &'static str);
pub struct Show;
pub struct ShowPrefixed(pub &'static str);
pub struct ShowShape;
pub struct ShowPrefixedShape(pub &'static str);
// Calls the given function.
pub struct Raw<T: Float, FUN: Fn(&NdArrayView<T>) + Send + Sync> {
pub(crate) raw: FUN,
pub(crate) phantom: PhantomData<T>,
}
impl<T: Float, FUN: Fn(&NdArrayView<T>) + Send + Sync> Hook<T> for Raw<T, FUN> {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
(self.raw)(arr)
}
}
impl<T: Float> Hook<T> for Print {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
println!("{}", self.0);
}
}
impl<T: Float> Hook<T> for Show {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
println!("{arr:?}");
}
}
impl<T: Float> Hook<T> for ShowPrefixed {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
println!("{} {:?}", self.0, arr);
}
}
impl<T: Float> Hook<T> for ShowShape {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
println!("{:?}", arr.shape());
}
}
impl<T: Float> Hook<T> for ShowPrefixedShape {
fn call(&self, arr: &crate::ndarray_ext::NdArrayView<T>) {
println!("{}\n{:?}", self.0, arr.shape());
}
}