sqrt.tch.rs

extern crate tch;  // https://github.com/LaurentMazare/tch-rs
use tch::{Tensor, no_grad};

fn sqrt(v: f64) -> f64 {
    let mut x = Tensor::from(v / 2.0).set_requires_grad(true);
    let lambda = 0.01;
    for _ in 0..100 {
        let y = &x * &x - v;
        let loss = &y * &y;
        x.zero_grad();
        loss.backward();
        println!("  y={} x={} dL/dx={}", y.double_value(&[]), x.double_value(&[]), x.grad().double_value(&[]));
        no_grad(|| {
            x -= x.grad() * lambda;
        });
    }
    x.double_value(&[])
}

fn main() {
    println!("sqrt 2 = {}", sqrt(2.0));
    println!("sqrt 3 = {}", sqrt(3.0));
    println!("sqrt 5 = {}", sqrt(5.0));
}