ferrite/
macros.rs

1#[macro_export]
2macro_rules! layer {
3  // The pattern: we match a single expression (`$x:expr`).
4  ($($x:tt)*) => {
5      // The expansion: we generate code that prints out the expression.
6      Box::new(Layer::$($x)*)
7  };
8}
9
10
11#[macro_export]
12macro_rules! grad_storage {
13  // The pattern: we match a single expression (`$x:expr`).
14  ($($x:tt)*) => {
15      // The expansion: we generate code that prints out the expression.
16      std::rc::Rc::new(std::cell::RefCell::new(Box::new($($x)*)))
17  };
18}
19
20#[macro_export]
21macro_rules! match_storage {
22  // Binary operation with two storage arguments
23  (binary $self:expr, $method:ident, $other:expr $(, $args:expr)*) => {
24    match ($self, $other) {
25      (Storage::Cpu(cpu_self), Storage::Cpu(cpu_other)) => {
26        Storage::Cpu(cpu_self.$method(cpu_other $(, $args)*))
27      }
28      _ => unimplemented!("Cross-device operations not supported"),
29    }
30  };
31
32  // Unary operation with single storage argument
33  (unary $self:expr, $method:ident $(, $args:expr)*) => {
34    match $self {
35      Storage::Cpu(cpu) => Storage::Cpu(cpu.$method($($args),*)),
36      _ => unimplemented!("Device not supported"),
37    }
38  };
39}
40
41#[macro_export]
42macro_rules! match_storage_assign {
43  // Binary operation with two storage arguments
44  (binary $self:expr, $method:ident, $other:expr $(, $args:expr)*) => {
45    match ($self, $other) {
46      (Storage::Cpu(cpu_self), Storage::Cpu(cpu_other)) => {
47        cpu_self.$method(cpu_other $(, $args)*)
48      }
49      _ => unimplemented!("Cross-device operations not supported"),
50    }
51  };
52
53  // Unary operation with single storage argument
54  (unary $self:expr, $method:ident $(, $args:expr)*) => {
55    match $self {
56      Storage::Cpu(cpu) => cpu.$method($($args,)*),
57      _ => unimplemented!("Device not supported"),
58    }
59  };
60}
61
62