Add the recip op + use it in stable-diffusion. (#331)

* Add the recip unary op.

* Fix the cuda kernel.

* Use the recip op in sigmoid.
This commit is contained in:
Laurent Mazare
2023-08-06 22:14:52 +02:00
committed by GitHub
parent 1c062bf06b
commit 166bfd5847
5 changed files with 26 additions and 5 deletions

View File

@ -1,7 +1,8 @@
use candle::{Result, Tensor};
use candle::{Device, Result, Tensor};
pub fn sigmoid(_: &Tensor) -> Result<Tensor> {
todo!()
pub fn sigmoid(xs: &Tensor) -> Result<Tensor> {
// TODO: Add sigmoid as binary ops.
(xs.neg()?.exp()? - 1.0)?.recip()
}
pub fn avg_pool2d(_: &Tensor) -> Result<Tensor> {
@ -16,6 +17,13 @@ pub fn upsample_nearest2d(_: &Tensor) -> Result<Tensor> {
todo!()
}
pub fn linspace(_: f64, _: f64, _: usize) -> Result<Tensor> {
todo!()
pub fn linspace(start: f64, stop: f64, steps: usize) -> Result<Tensor> {
if steps < 1 {
candle::bail!("cannot use linspace with steps {steps} <= 1")
}
let delta = (stop - start) / (steps - 1) as f64;
let vs = (0..steps)
.map(|step| start + step as f64 * delta)
.collect::<Vec<_>>();
Tensor::from_vec(vs, steps, &Device::Cpu)
}