Add BackendRouter to handle multiple backends

[laggui]

Needs more tests :)

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Closes #2276

Changes

Introduces a new BackendRouter responsible for forwarding tensor operations to the appropriate backend, given multiple backends.

This is achieved with the help of the intermediate representation defined for ReprBackend and the tensor/ops descriptions.

Testing

Modified the ag-news-train text classification example to run on cuda + wgpu. Also have a MWE:

use burn::{
    backend::{cuda_jit::CudaDevice, wgpu::WgpuDevice, CudaJit, Wgpu},
    tensor::Tensor,
};
use burn_router::{BackendRouter, ByteBridge, DirectChannel, MultiDevice2};

fn main() {
    type DirectByteChannel<Backends> = DirectChannel<Backends, ByteBridge<Backends>>;

    type DualBackend = BackendRouter<DirectByteChannel<(CudaJit, Wgpu)>>;

    let device2 = WgpuDevice::Cpu;
    let device1 = CudaDevice::new(0);

    // TODO: this is wack.. how to automatically implement From<B1::Device1> for MultiDevice2?
    let multi_device1 = MultiDevice2::Device1(device1);
    let multi_device2 = MultiDevice2::Device2(device2);
    let tensor1 = Tensor::<DualBackend, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &multi_device1);
    let tensor2 = Tensor::<DualBackend, 1>::from_floats([5.0, 6.0, 7.0, 8.0], &multi_device2);

    println!("Tensor 1:\n{tensor1}");
    println!("Tensor 2:\n{tensor2}");

    let tensor1 = tensor1.to_device(&multi_device2);

    let output = tensor1.add(tensor2);

    println!("Result:\n{output}");
}

Tensor 1:
Tensor {
  data:
[1.0, 2.0, 3.0, 4.0],
  shape:  [4],
  device:  Device1(CudaDevice { index: 0 }),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}
Tensor 2:
Tensor {
  data:
[5.0, 6.0, 7.0, 8.0],
  shape:  [4],
  device:  Device2(Cpu),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}
Result:
Tensor {
  data:
[6.0, 8.0, 10.0, 12.0],
  shape:  [4],
  device:  Device2(Cpu),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}

[codecov]

Codecov Report

Attention: Patch coverage is 0.16560% with 4823 lines in your changes missing coverage. Please review.

Project coverage is 81.10%. Comparing base (c98b689) to head (0452d2c).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-router/src/ops/op_float.rs	0.00%	1299 Missing ⚠️
crates/burn-router/src/ops/op_int.rs	0.00%	1064 Missing ⚠️
crates/burn-router/src/runner.rs	0.00%	1001 Missing ⚠️
crates/burn-router/src/ops/op_module.rs	0.00%	768 Missing ⚠️
crates/burn-router/src/ops/op_bool.rs	0.00%	216 Missing ⚠️
crates/burn-router/src/bridge/byte.rs	0.00%	105 Missing ⚠️
crates/burn-router/src/channel/direct.rs	0.00%	99 Missing ⚠️
crates/burn-router/src/tensor.rs	0.00%	80 Missing ⚠️
crates/burn-router/src/ops/op_qfloat.rs	0.00%	54 Missing ⚠️
crates/burn-router/src/client/base.rs	0.00%	44 Missing ⚠️
... and 5 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2353      +/-   ##
==========================================
- Coverage   85.23%   81.10%   -4.14%     
==========================================
  Files         770      782      +12     
  Lines       98976   103525    +4549     
==========================================
- Hits        84358    83959     -399     
- Misses      14618    19566    +4948     

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