Repeat ops autodiff & fusion + fix autodiff ones & zeros

[louisfd]

I implemented the backward pass for repeat, enabling it to use the jit kernel of repeat instead of the fallback implementation with slice assign.
I also added fusion support for repeat

Also fixed ones and zeros in autodiff which were mistakingly cpu-bound

[laggui]

Was curious to try the difference with the default implementation..

use std::time::Instant;

use burn::{
    backend::{wgpu::AutoGraphicsApi, Autodiff, Wgpu},
    tensor::backend::Backend,
    tensor::Tensor,
};

fn main() {
    type B = Autodiff<Wgpu<AutoGraphicsApi, f32, i32>>;
    let device = Default::default();
    let now = Instant::now();
    let t = Tensor::<B, 1>::from_floats([1.0], &device);
    let _rep = t.repeat(0, 1024 * 1024);

    B::sync(&device);
    let elapsed = now.elapsed().as_millis();
    println!("Took {} ms", elapsed);
}

With the default impl on main: Took 36066 ms
With your fix: Took 387 ms (multiple attempts varied between 300-1000ms)

... a timely fix 😄

[codecov]

Codecov Report

Attention: Patch coverage is 93.78238% with 12 lines in your changes are missing coverage. Please review.

Project coverage is 86.36%. Comparing base (a880d2c) to head (8783cd8).
Report is 2 commits behind head on main.

Files	Patch %	Lines
crates/burn-autodiff/src/ops/bool_tensor.rs	0.00%	7 Missing ⚠️
crates/burn-autodiff/src/ops/tensor.rs	90.56%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1600      +/-   ##
==========================================
- Coverage   86.39%   86.36%   -0.03%     
==========================================
  Files         688      689       +1     
  Lines       78675    78866     +191     
==========================================
+ Hits        67973    68116     +143     
- Misses      10702    10750      +48     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.