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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
use smallvec::smallvec;
use crate::{Element, MapOperation, Recordable, Shape, Tensor};
use super::{Cotangents, Operation, Reads, unary};
/// A unary elementwise transcendental of a value: one node kind
/// carrying the [`MapOperation`] it applies.
///
/// The IR and the backend map seam share this vocabulary on purpose:
/// `tanh` recorded here and `tanh` offered to the backend chain are
/// the same instruction, so adding a transcendental is a
/// `MapOperation` variant and the arms below — never a new
/// `Op` variant. Everything op-specific (the printed name, the
/// read set, the derivative) dispatches on `op`; the shape behavior
/// is shared, since a map always keeps its operand's shape.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Map {
pub(crate) op: MapOperation,
}
impl Map {
/// Returns the arity: one operand.
pub(crate) fn arity(&self) -> usize {
1
}
/// Returns the read set of the derivative rules below, per
/// operation: `Exp`, `Sqrt`, `Tanh`, and `Expm1` reuse their own
/// output; `Ln`, `Sin`, `Cos`, `Log1p`, and `Erf` read their
/// operand; `ErfDerivative` reads both. The set is deliberately
/// not uniform — a shared one would retain buffers liveness does
/// not need.
pub(crate) fn reads(&self) -> Reads {
match self.op {
MapOperation::Exp | MapOperation::Sqrt | MapOperation::Tanh | MapOperation::Expm1 => {
Reads {
operands: [false, false],
output: true,
}
}
MapOperation::Ln
| MapOperation::Sin
| MapOperation::Cos
| MapOperation::Log1p
| MapOperation::Erf => Reads {
operands: [true, false],
output: false,
},
MapOperation::ErfDerivative => Reads {
operands: [true, false],
output: true,
},
}
}
/// Infers the shape of the result: the operand's shape.
pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
unary(operands).clone()
}
}
impl Map {
pub(crate) fn forward<E: Element>(&self, operands: &[&Tensor<E>]) -> Tensor<E> {
let &operand = unary(operands);
match self.op {
MapOperation::Exp => operand.exp(),
MapOperation::Ln => operand.ln(),
MapOperation::Sqrt => operand.sqrt(),
MapOperation::Tanh => operand.tanh(),
MapOperation::Sin => operand.sin(),
MapOperation::Cos => operand.cos(),
MapOperation::Log1p => operand.log1p(),
MapOperation::Expm1 => operand.expm1(),
MapOperation::Erf => operand.erf(),
MapOperation::ErfDerivative => operand.erf_derivative(),
}
}
}
impl<Rule: Recordable> Operation<Rule> for Map {
fn backward(&self, operands: &[&Rule], output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
let cotangent = match self.op {
// The derivative of `e^x` is `e^x` itself: the canonical
// case of reusing the node's own output.
MapOperation::Exp => gradient.clone() * output.clone(),
// The derivative of `ln(x)` is `1 / x`; gradients inherit
// the payload's logarithm and division semantics outside
// the positive domain.
MapOperation::Ln => {
let &operand = unary(operands);
gradient.clone() / operand.clone()
}
// The derivative of `sqrt(x)` is `1 / (2 * sqrt(x))` —
// no generic literal `2` exists, so the doubling is
// `output + output`.
MapOperation::Sqrt => gradient.clone() / (output.clone() + output.clone()),
// The derivative of `tanh(x)` is `1 - tanh(x)^2`: one
// minus the square of the node's own output.
MapOperation::Tanh => {
gradient.clone() * (output.one_like() - output.clone() * output.clone())
}
// The derivative of `sin(x)` is `cos(x)`: the pair closes
// over itself, which is why the two ship together.
MapOperation::Sin => {
let &operand = unary(operands);
gradient.clone() * operand.cos()
}
// The derivative of `cos(x)` is `-sin(x)`.
MapOperation::Cos => {
let &operand = unary(operands);
-(gradient.clone() * operand.sin())
}
// The derivative of `ln(1 + x)` is `1 / (1 + x)`; the
// fused accuracy is a forward property, and the rule's
// own `1 + x` is safe — no cancellation hides in an
// addition this side of the logarithm.
MapOperation::Log1p => {
let &operand = unary(operands);
gradient.clone() / (operand.one_like() + operand.clone())
}
// The derivative of `e^x - 1` is `e^x`: the node's own
// output plus one, mirroring `Exp`'s output reuse.
MapOperation::Expm1 => gradient.clone() * (output.clone() + output.one_like()),
// The derivative of `erf(x)` is the scaled Gaussian — its
// own operation, so the rule mints no constant.
MapOperation::Erf => {
let &operand = unary(operands);
gradient.clone() * operand.erf_derivative()
}
// The derivative of the scaled Gaussian is `-2x` times
// itself: the node's own output, doubled operand, negated.
MapOperation::ErfDerivative => {
let &operand = unary(operands);
-(gradient.clone() * (operand.clone() + operand.clone()) * output.clone())
}
};
smallvec![Some(cotangent)]
}
}
#[cfg(test)]
#[path = "tests/map_tests.rs"]
mod tests;