为什么要转

model = Mlp().eval()
with torch.no_grad():
  model.fc1.weight.copy_(w(HID, IN, math.sin, 0.5))
  model.fc1.bias.copy_(torch.tensor([0.01 * o for o in range(HID)]))
  model.fc2.weight.copy_(w(OUT, HID, math.cos, 0.5))
  model.fc2.bias.copy_(torch.tensor([-0.02 * o for o in range(OUT)]))

x = torch.tensor([[0.1 * (r * IN + c) - 0.5 for c in range(IN)] for r in range(4)])
with torch.inference_mode():
  logits = model(x)
json.dump({"logits": logits.tolist()}, sys.stdout)

fn main() {
  let device = Device::flex();
  let mut model = Mlp {
    fc1: LinearConfig::new(IN, HID).init(&device),
    fc2: LinearConfig::new(HID, OUT).init(&device),
    relu: Relu::new(),
  };
  model.fc1.weight = Param::from_tensor(weight(HID, IN, f64::sin, &device));
  model.fc1.bias = Some(Param::from_tensor(bias(HID, 0.01, &device)));
  model.fc2.weight = Param::from_tensor(weight(OUT, HID, f64::cos, &device));
  model.fc2.bias = Some(Param::from_tensor(bias(OUT, -0.02, &device)));

  let xs: Vec<f32> = (0..4 * IN).map(|k| 0.1 * k as f32 - 0.5).collect();
  let x = Tensor::<2>::from_data(TensorData::new(xs, [4, IN]), &device);
  let logits = model.forward(x);
  let v: Vec<f32> = match logits.into_data().try_to_vec::<f32>() {
    Ok(v) => v,
    Err(e) => {
      eprintln!("read back failed: {e:?}");
      std::process::exit(1);
    }
  };
  let rows: Vec<&[f32]> = v.chunks(OUT).collect();
  println!("{}", json!({ "logits": rows }));
}

[dependencies]
burn = { version = "=0.22.1", default-features = false, features = ["std", "flex"] }
serde_json = "1"

[profile.release]
strip = true

{"logits": [[-0.09422215819358826, -0.12721742689609528, 0.08542244881391525], [0.0261688269674778, -0.12934733927249908, -0.0343487448990345], [0.056371480226516724, -0.20086252689361572, -0.04374047741293907], [0.11677984893321991, -0.35583826899528503, -0.059050895273685455]]}

{"logits":[[-0.09422216564416885,-0.12721742689609528,0.08542244881391525],[0.0261688232421875,-0.12934733927249908,-0.0343487374484539],[0.056371480226516724,-0.20086251199245453,-0.043740492314100266],[0.1167798638343811,-0.35583826899528503,-0.05905091017484665]]}

logits                   n= 12  max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK

Requires-Dist: cuda-toolkit[cublas,cudart,cufft,cufile,cupti,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==13.2.2; platform_system == "Linux"
Requires-Dist: nvidia-cudnn-cu13==9.26.0.51; platform_system == "Linux"
...
Requires-Dist: triton~=3.9.0; platform_system == "Linux" and python_version < "3.16"

for name, cmd in cmds.items():
  subprocess.run(cmd, stdout=subprocess.DEVNULL, check=True)  # 预热文件缓存
  ts = []
  for _ in range(N):
    t0 = time.perf_counter()
    subprocess.run(cmd, stdout=subprocess.DEVNULL, check=True)
    ts.append((time.perf_counter() - t0) * 1000)
  print(f"{name:30} n={N}  min={min(ts):8.1f} ms  median={statistics.median(ts):8.1f} ms")

python -I -c pass              n=30  min=    13.2 ms  median=    13.7 ms
python -I -c 'import torch'    n=30  min=   650.0 ms  median=   661.2 ms
python -I intro_infer.py       n=30  min=   653.8 ms  median=   663.9 ms
intro-infer (release)          n=30  min=     2.6 ms  median=     2.9 ms
intro-infer-default (release)  n=30  min=     2.6 ms  median=     2.8 ms
intro-infer +metal (release)   n=30  min=     5.6 ms  median=     5.9 ms

import torch 552.9 ms, build 0.43 ms, forward#1 0.133 ms, forward#2 0.016 ms, threads 8
import torch 542.8 ms, build 0.40 ms, forward#1 0.135 ms, forward#2 0.015 ms, threads 8
import torch 546.9 ms, build 0.38 ms, forward#1 0.127 ms, forward#2 0.017 ms, threads 8
import torch 547.1 ms, build 0.40 ms, forward#1 0.128 ms, forward#2 0.015 ms, threads 8
import torch 546.6 ms, build 0.38 ms, forward#1 0.125 ms, forward#2 0.018 ms, threads 8

a = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
print("cumsum(1) =", a.cumsum(1).flatten().tolist())
print("einsum ij,j->i =", torch.einsum("ij,j->i", a, torch.tensor([10.0, 20.0])).tolist())
print("svd S =", torch.linalg.svd(a).S.tolist())
g = a.clone().requires_grad_()
torch.linalg.svd(g).S.sum().backward()
print("svd grad =", g.grad.flatten().tolist())

use burn::linalg::svd;
use burn::tensor::{Device, Tensor};

fn main() {
  let device = Device::flex();
  let a = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device);
  println!("cumsum(1) = {:?}", a.clone().cumsum(1).into_data().try_to_vec::<f32>());
  let v = Tensor::<1>::from_floats([10.0, 20.0], &device);
  let mv = Tensor::<1>::einsum("ij,j->i", [a.clone().into(), v.into()]);
  println!("einsum ij,j->i = {:?}", mv.into_data().try_to_vec::<f32>());
  let (u, s, vt) = svd::<2, 1>(a.clone(), 10);
  println!("svd U {:?} Vt {:?}", u.dims(), vt.dims());
  println!("svd S = {:?}", s.into_data().try_to_vec::<f32>());

  if std::env::args().any(|x| x == "--grad") {
    let ad = Device::flex().autodiff();
    let g = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &ad).require_grad();
    let (_u, s, _vt) = svd::<2, 1>(g, 10);
    println!("{:?}", s.into_data().try_to_vec::<f32>());
  }
}

burn = { version = "=0.22.1", default-features = false, features = ["std", "flex", "autodiff", "linalg"] }

cumsum(1) = [1.0, 3.0, 3.0, 7.0, 5.0, 11.0]
einsum ij,j->i = [50.0, 110.0, 170.0]
svd S = [9.525519371032715, 0.5143014788627625]
svd grad = [-0.5510035753250122, 0.7278250455856323, 0.13615891337394714, 0.5610649585723877, 0.8233200311660767, 0.3943061828613281]

cumsum(1) = Ok([1.0, 3.0, 3.0, 7.0, 5.0, 11.0])
einsum ij,j->i = Ok([50.0, 110.0, 170.0])
svd U [3, 2] Vt [2, 2]
svd S = Ok([9.5255165, 0.5143005])

thread 'main' (8509250) panicked at .../burn-linalg-0.22.1/src/functions/svd.rs:98:9:
linalg::svd: gradients are not implemented; detach the input tensor first

t = torch.tensor([3, 1, 3, 2])
print("unique =", torch.unique(t).tolist())
print("searchsorted =", torch.searchsorted(torch.tensor([0.0, 1.0, 2.0]), torch.tensor([0.5, 1.5])).tolist())
print("bincount =", torch.bincount(t).tolist())
print("inv =", torch.linalg.inv(torch.tensor([[2.0, 0.0], [0.0, 4.0]])).tolist())

unique = [1, 2, 3]
searchsorted = [1, 2]
bincount = [0, 1, 1, 2]
inv = [[0.5, 0.0], [0.0, 0.25]]

use burn::tensor::{Device, Int, Tensor};

fn main() {
  let device = Device::flex();
  let t = Tensor::<1, Int>::from_ints([3, 1, 3, 2], &device);
  let _u = t.clone().unique();
  let edges = Tensor::<1>::from_floats([0.0, 1.0, 2.0], &device);
  let _i = edges.searchsorted(Tensor::<1>::from_floats([0.5, 1.5], &device));
  let _c = t.bincount();
  let m = Tensor::<2>::from_floats([[2.0, 0.0], [0.0, 4.0]], &device);
  let _inv = burn::linalg::inv(m);
}

error[E0425]: cannot find function `inv` in module `burn::linalg`
  --> src/bin/missing.rs:13:28
   |
13 |   let _inv = burn::linalg::inv(m);
   |                            ^^^ not found in `burn::linalg`

error[E0599]: no method named `unique` found for struct `Tensor<D, K>` in the current scope
 --> src/bin/missing.rs:8:22
  |
8 |   let _u = t.clone().unique();
  |                      ^^^^^^ method not found in `Tensor<1, burn::tensor::Int>`

error[E0599]: no method named `searchsorted` found for struct `Tensor<D, K>` in the current scope
  --> src/bin/missing.rs:10:18
   |
10 |   let _i = edges.searchsorted(Tensor::<1>::from_floats([0.5, 1.5], &device));
   |                  ^^^^^^^^^^^^
   |
help: there is a method `argsort` with a similar name
...
error[E0599]: no method named `bincount` found for struct `Tensor<D, K>` in the current scope
  --> src/bin/missing.rs:11:14
   |
11 |   let _c = t.bincount();
   |              ^^^^^^^^ method not found in `Tensor<1, burn::tensor::Int>`