-
Notifications
You must be signed in to change notification settings - Fork 8
Description
Reproducer:
#![feature(autodiff)]
use std::autodiff::autodiff;
#[autodiff(sin_vec, Reverse, Duplicated, Const, Active)]
fn cos_vec(x: &mut Vec<f32>, y: usize) -> f32 {
// uses enum internally and breaks
for i in 0..y {
x.push(i as f32);
}
let res = x.into_iter().collect::<Vec<&mut f32>>();
*res[0]
}
fn main() {
let mut x = vec![1.0, 1.0, 1.0];
let mut d_x = vec![0.0; 3];
sin_vec(&mut x, &mut d_x, 4, 1.0);
dbg!(&d_x, &x);
}
RUSTFLAGS="-Z autodiff=LooseTypes,Print" cargo +enzyme run
Demangled function names:
https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=2fdf75b03825f5f63f353ff1122f3e0f
Demangled function names are:
alloc::raw_vec::finish_grow::hce1ff89b620ba31c
alloc::raw_vec::RawVecInner<A>::grow_amortized::h6d11486032b72786
alloc::raw_vec::RawVecInner<A>::reserve::do_reserve_and_handle::hf69592d80ba560b8
alloc::vec::Vec<T,A>::reserve::hfbf4bb7fa5035d73
alloc::vec::Vec<T,A>::extend_trusted::h9fccf647fedf83c0
I'd say reserve is a reasonable function to provide a custom-derivative for @wsmoses , what would we need to register?
Augmented primal is just a reserve for the same additional on the shadow, and do nothing for the reverse?
For context, that's the definition:
pub struct Vec<T, A: Allocator = Global> {
buf: RawVec<T, A>,
len: usize,
}
pub fn reserve(&mut self, additional: usize) {
self.buf.reserve(self.len, additional);
}
This example and some playing around in release mode indicates we should also have a custom rule for alloc::raw_vec::RawVecInner<A>::reserve::do_reserve_and_handle
https://play.rust-lang.org/?version=nightly&mode=release&edition=2021&gist=3d81e4ea9dcd7abd8acace3d9e40685c