You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
@@ -154,8 +154,6 @@ The implementation of @ref stan::math::make_holder is [here](https://github.com/
154
154
155
155
### Move Semantics
156
156
157
-
In general, Stan Math does not use move semantics very often.
158
-
This is because of our arena allocator.
159
157
Move semantics generally work as
160
158
161
159
```cpp
@@ -179,6 +177,96 @@ We can see in the above that the standard style of a move (the constructor takin
179
177
But in Stan, particularly for reverse mode, we need to keep memory around even if it's only temporary for when we call the gradient calculations in the reverse pass.
180
178
And since memory for reverse mode is stored in our arena allocator no copying happens in the first place.
181
179
180
+
Functions for Stan Math's reverse mode autodiff should use [_perfect forwarding_](https://drewcampbell92.medium.com/understanding-move-semantics-and-perfect-forwarding-part-3-65575d523ff8) arguments. Perfect forwarding arguments use a template parameter wit no attributes such as `const` and `volatile` and have a double ampersand `&&` next to them.
181
+
182
+
```c++
183
+
template <typename T>
184
+
auto my_function(T&& x) {
185
+
return my_other_function(std::forward<T>(x));
186
+
}
187
+
```
188
+
189
+
The `std::forward<T>` in the in the code above tells the compiler that if `T` is deduced to be an rvalue reference (such as `Eigen::MatrixXd&&`), then it should be moved to `my_other_function`, where there it can possibly use another objects move constructor to reuse memory.
190
+
A perfect forwarding argument of a function accepts any reference type as its input argument.
191
+
The above signature is equivalent to writing out several functions with different reference types
192
+
193
+
```c++
194
+
// Accepts a plain lvalue reference
195
+
automy_function(Eigen::MatrixXd& x) {
196
+
return my_other_function(x);
197
+
}
198
+
// Accepts a const lvalue reference
199
+
auto my_function(const Eigen::MatrixXd& x) {
200
+
return my_other_function(x);
201
+
}
202
+
// Accepts an rvalue reference
203
+
auto my_function(Eigen::MatrixXd&& x) {
204
+
return my_other_function(std::move(x));
205
+
}
206
+
// Accepts a const rvalue reference
207
+
auto my_function(const Eigen::MatrixXd&& x) {
208
+
return my_other_function(std::move(x));
209
+
}
210
+
```
211
+
212
+
In Stan, perfect forwarding is used in reverse mode functions which can accept an Eigen matrix type.
The signature for this function has a template `T` that is required to be an Eigen type with a `value_type` that is a `var` type.
235
+
The template parameter `T` is then used in the signature as an perfect forwarding argument.
236
+
237
+
```c++
238
+
// Store `x` on the arena
239
+
arena_t<T> x_arena(std::forward<T>(x));
240
+
```
241
+
242
+
The input is stored in the arena, which is where the perfect forwarding magic actually occurs.
243
+
If `T` is an lvalue type such as `Eigen::MatrixXd&` then `arena_matrix` will use it's copy constructor, creating new memory in Stan's arena allocator and then copying the values of `x` into that memory.
244
+
But if `T` was a temporary rvalue type such as `Eigen::MatrixXd&&`, then the `arena_matrix` class will use it's move constructor to place the temporary matrix in Stan's `var_alloc_stack_`.
245
+
The `var_alloc_stick_` is used to hold objects that were created outside of the arena allocator but need to be deleted when the arena allocator is cleared.
246
+
This allows the `arena_matrix` to reuse the memory from the temporary matrix. Then the matrix will be deleted once arena allocator requests memory to be cleared.
The rest of this code follows the standard format for the rest of Stan Math's reverse mode that accepts Eigen types as input.
264
+
The `reverse_pass_callback` function accepts a lambda as input and places the lambda in Stan's callback stack to be called later when `grad()` is called by the user.
265
+
Since `arena_matrix` types only store a pointer to memory allocated elsewhere they are copied into the lambda.
266
+
The body of the lambda holds the gradient calculation needed for the reverse mode pass.
267
+
268
+
Then finally `ret`, the `arena_matrix` type is returned by the function.
269
+
182
270
When working with arithmetic types, keep in mind that moving Scalars is often less optimal than simply taking their copy.
183
271
For instance, Stan's `var` type uses the pointer to implementation (PIMPL) pattern, so it simply holds a pointer of size 8 bytes.
184
272
A `double` is also 8 bytes which just so happens to fit exactly in a [word](https://en.wikipedia.org/wiki/Word_(computer_architecture)) of most modern CPUs with at least 64-byte cache lines.
@@ -190,6 +278,45 @@ The general rules to follow for passing values to a function are:
190
278
2. If you are writing a function for reverse mode, pass values by `const&`
191
279
3. In prim, if you are confident and working with larger types, use perfect forwarding to pass values that can be moved from. Otherwise simply pass values by `const&`.
192
280
281
+
### Using auto is Dangerous With Eigen Matrix Functions in Reverse Mode
282
+
283
+
The use of auto with the Stan Math library should be used with care, like in [Eigen](https://eigen.tuxfamily.org/dox/TopicPitfalls.html).
284
+
Along with the cautions mentioned in the Eigen docs, there are also memory considerations when using reverse mode automatic differentiation.
285
+
When returning from a function in the Stan Math library with an Eigen matrix output with a scalar `var` type, the actual returned type will often be an `arena_matrix<Eigen::Matrix<...>>`.
286
+
The `arena_matrix` class is an Eigen matrix where the underlying array of memory is located in Stan's memory arena.
287
+
The `arena_matrix` that is returned by Math functions is normally the same one resting in the callback used to calculate gradients in the reverse pass.
288
+
Directly changing the elements of this matrix would also change the memory the reverse pass callback sees which would result in incorrect calculations.
289
+
290
+
The simple solution to this is that when you use a math library function that returns a matrix and then want to assign to any of the individual elements of the matrix, assign to an actual Eigen matrix type instead of using auto.
291
+
In the below example, we see the first case which uses auto and will change the memory of the `arena_matrix` returned in the callback for multiply's reverse mode.
292
+
Directly below it is the safe version, which just directly assigns to an Eigen matrix type and is safe to do element insertion into.
293
+
294
+
```c++
295
+
Eigen::Matrix<var, -1, 1> y;
296
+
Eigen::Matrix<var, -1, -1> X;
297
+
// Bad!! Will change memory used by reverse pass callback within multiply!
298
+
auto mu = multiply(X, y);
299
+
mu(4) = 1.0;
300
+
// Good! Will not change memory used by reverse pass callback within multiply
The reason we do this is for cases where function returns are passed to other functions.
306
+
An `arena_matrix` will always make a shallow copy when being constructed from another `arena_matrix`, which lets the functions avoid unnecessary copies.
307
+
308
+
```c++
309
+
Eigen::Matrix<var, -1, 1> y1;
310
+
Eigen::Matrix<var, -1, -1> X1;
311
+
Eigen::Matrix<var, -1, 1> y2;
312
+
Eigen::Matrix<var, -1, -1> X2;
313
+
auto mu1 = multiply(X1, y1);
314
+
auto mu2 = multiply(X2, y2);
315
+
// Inputs not copied in this case!
316
+
auto z = add(mu1, mu2);
317
+
```
318
+
319
+
193
320
### Passing variables that need destructors called after the reverse pass (`make_chainable_ptr`)
194
321
195
322
When possible, non-arena variables should be copied to the arena to be used in the reverse pass.
@@ -242,22 +369,17 @@ grad();
242
369
```
243
370
244
371
Now `res` is `innocent_return` and we've changed one of the elements of `innocent_return`, but that is also going to change the element of `res` which is being used in our reverse pass callback!
245
-
The answer for this is simple but sadly requires a copy.
Care must be taken by end users of Stan Math by using `auto` with caution.
374
+
When a user wishes to manipulate the coefficients of a matrix that is a return from a function in Stan Math, they should assign the matrix to a plain Eigen type.
258
375
259
-
we make a deep copy of the return whose inner `vari` will not be the same, but the `var` will produce a new copy of the pointer to the `vari`.
260
-
Now the user code above will be protected, and it is safe for them to assign to individual elements of the `auto` returned matrix.
376
+
```c++
377
+
Eigen::Matrix<var, -1, 1> x = Eigen::Matrix<double, -1, 1>::Random(5);
0 commit comments