Automated conversion of list arguments (2). (#682)

* Add the list conversion traits and use them in the code generation.

* Automated fix of all clippy lints.

* Simplify the calls to sum_dim_intlist.

CHANGELOG.md

@@ -10,6 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   weight by default, [679](https://github.com/LaurentMazare/tch-rs/pull/679).
 - None can be used for nullable scalar types,
   [680](https://github.com/LaurentMazare/tch-rs/pull/680).
+- Automated conversion of list arguments: all the generated functions that take
+  as input a slice of int or float can now be used directly with int values or
+  fixed length arrays [682](https://github.com/LaurentMazare/tch-rs/pull/682).
 
 ## v0.11.0 - 2023-03-20
 ### Added

Cargo.toml

@@ -22,7 +22,7 @@ libc = "0.2.0"
 ndarray = "0.15"
 rand = "0.8"
 thiserror = "1"
-torch-sys = { version = "0.11.0", path = "torch-sys" }
+torch-sys = { version = "0.12.0", path = "torch-sys" }
 zip = "0.6"
 half = "2"
 safetensors = "0.3.0"

examples/basics.rs

@@ -20,7 +20,7 @@ fn main() {
     let device = tch::Device::cuda_if_available();
     let t = Tensor::of_slice(&[3, 1, 4, 1, 5]).to(device);
     t.print();
-    let t = Tensor::randn(&[5, 4], kind::FLOAT_CPU);
+    let t = Tensor::randn([5, 4], kind::FLOAT_CPU);
     t.print();
     (&t + 1.5).print();
     (&t + 2.5).print();

examples/char-rnn/main.rs

@@ -23,7 +23,7 @@ fn sample(data: &TextData, lstm: &LSTM, linear: &Linear, device: Device) -> Stri
     let mut last_label = 0i64;
     let mut result = String::new();
     for _index in 0..SAMPLING_LEN {
-        let input = Tensor::zeros(&[1, labels], (Kind::Float, device));
+        let input = Tensor::zeros([1, labels], (Kind::Float, device));
         let _ = input.narrow(1, last_label, 1).fill_(1.0);
         state = lstm.step(&input, &state);
         let sampled_y = linear

examples/gan/main.rs

@@ -86,11 +86,11 @@ pub fn main() -> Result<()> {
     let train_size = images.size()[0];
 
     let random_batch_images = || {
-        let index = Tensor::randint(train_size, &[BATCH_SIZE], kind::INT64_CPU);
+        let index = Tensor::randint(train_size, [BATCH_SIZE], kind::INT64_CPU);
         images.index_select(0, &index).to_device(device).to_kind(Kind::Float) / 127.5 - 1.
     };
     let rand_latent = || {
-        (Tensor::rand(&[BATCH_SIZE, LATENT_DIM, 1, 1], kind::FLOAT_CPU) * 2.0 - 1.0)
+        (Tensor::rand([BATCH_SIZE, LATENT_DIM, 1, 1], kind::FLOAT_CPU) * 2.0 - 1.0)
             .to_device(device)
     };
 

examples/jit-trace/main.rs

@@ -54,7 +54,7 @@ fn main() -> Result<()> {
     let model = CModule::create_by_tracing(
         "MyModule",
         "forward",
-        &[Tensor::zeros(&[784], FLOAT_CUDA)],
+        &[Tensor::zeros([784], FLOAT_CUDA)],
         &mut closure,
     )?;
     model.save("model.pt")?;

examples/min-gpt/main.rs

@@ -69,7 +69,7 @@ fn causal_self_attention(p: &nn::Path, cfg: Config) -> impl ModuleT {
     let value = linear(p / "value", cfg.n_embd, cfg.n_embd);
     let proj = linear(p / "proj", cfg.n_embd, cfg.n_embd);
     let mask_init =
-        Tensor::ones(&[cfg.block_size, cfg.block_size], (Kind::Float, p.device())).tril(0);
+        Tensor::ones([cfg.block_size, cfg.block_size], (Kind::Float, p.device())).tril(0);
     let mask_init = mask_init.view([1, 1, cfg.block_size, cfg.block_size]);
     // let mask = p.var_copy("mask", &mask_init);
     let mask = mask_init;
@@ -180,7 +180,7 @@ pub fn main() -> Result<()> {
                     idx += 1;
                     if idx % 10000 == 0 {
                         println!("Epoch: {}   loss: {:5.3}", epoch, sum_loss / cnt_loss);
-                        let input = Tensor::zeros(&[1, BLOCK_SIZE], (Kind::Int64, device));
+                        let input = Tensor::zeros([1, BLOCK_SIZE], (Kind::Int64, device));
                         println!("Sample: {}", sample(&data, &gpt, input));
                         if let Err(err) = vs.save(format!("gpt{idx}.ot")) {
                             println!("error while saving {err}");
@@ -194,7 +194,7 @@ pub fn main() -> Result<()> {
         "predict" => {
             vs.load(args[2].as_str())?;
             let seqstart = args[3].as_str();
-            let input = Tensor::zeros(&[1, BLOCK_SIZE], (Kind::Int64, device));
+            let input = Tensor::zeros([1, BLOCK_SIZE], (Kind::Int64, device));
             for (idx, c) in seqstart.chars().rev().enumerate() {
                 let idx = idx as i64;
                 if idx >= BLOCK_SIZE {

examples/mnist/mnist_linear.rs

@@ -12,8 +12,8 @@ pub fn run() -> Result<()> {
     println!("train-labels: {:?}", m.train_labels.size());
     println!("test-images: {:?}", m.test_images.size());
     println!("test-labels: {:?}", m.test_labels.size());
-    let mut ws = Tensor::zeros(&[IMAGE_DIM, LABELS], kind::FLOAT_CPU).set_requires_grad(true);
-    let mut bs = Tensor::zeros(&[LABELS], kind::FLOAT_CPU).set_requires_grad(true);
+    let mut ws = Tensor::zeros([IMAGE_DIM, LABELS], kind::FLOAT_CPU).set_requires_grad(true);
+    let mut bs = Tensor::zeros([LABELS], kind::FLOAT_CPU).set_requires_grad(true);
     for epoch in 1..200 {
         let logits = m.train_images.mm(&ws) + &bs;
         let loss = logits.log_softmax(-1, Kind::Float).nll_loss(&m.train_labels);

examples/reinforcement-learning/a2c.rs

@@ -46,7 +46,7 @@ struct FrameStack {
 
 impl FrameStack {
     fn new(nprocs: i64, nstack: i64) -> FrameStack {
-        FrameStack { data: Tensor::zeros(&[nprocs, nstack, 84, 84], FLOAT_CPU), nprocs, nstack }
+        FrameStack { data: Tensor::zeros([nprocs, nstack, 84, 84], FLOAT_CPU), nprocs, nstack }
     }
 
     fn update<'a>(&'a mut self, img: &Tensor, masks: Option<&Tensor>) -> &'a Tensor {
@@ -73,19 +73,19 @@ pub fn train() -> cpython::PyResult<()> {
     let model = model(&vs.root(), env.action_space());
     let mut opt = nn::Adam::default().build(&vs, 1e-4).unwrap();
 
-    let mut sum_rewards = Tensor::zeros(&[NPROCS], FLOAT_CPU);
+    let mut sum_rewards = Tensor::zeros([NPROCS], FLOAT_CPU);
     let mut total_rewards = 0f64;
     let mut total_episodes = 0f64;
 
     let mut frame_stack = FrameStack::new(NPROCS, NSTACK);
     let _ = frame_stack.update(&env.reset()?, None);
-    let s_states = Tensor::zeros(&[NSTEPS + 1, NPROCS, NSTACK, 84, 84], FLOAT_CPU);
+    let s_states = Tensor::zeros([NSTEPS + 1, NPROCS, NSTACK, 84, 84], FLOAT_CPU);
     for update_index in 0..UPDATES {
         s_states.get(0).copy_(&s_states.get(-1));
-        let s_values = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
-        let s_rewards = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
-        let s_actions = Tensor::zeros(&[NSTEPS, NPROCS], INT64_CPU);
-        let s_masks = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
+        let s_values = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
+        let s_rewards = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
+        let s_actions = Tensor::zeros([NSTEPS, NPROCS], INT64_CPU);
+        let s_masks = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
         for s in 0..NSTEPS {
             let (critic, actor) = tch::no_grad(|| model(&s_states.get(s)));
             let probs = actor.softmax(-1, Float);
@@ -106,7 +106,7 @@ pub fn train() -> cpython::PyResult<()> {
             s_masks.get(s).copy_(&masks);
         }
         let s_returns = {
-            let r = Tensor::zeros(&[NSTEPS + 1, NPROCS], FLOAT_CPU);
+            let r = Tensor::zeros([NSTEPS + 1, NPROCS], FLOAT_CPU);
             let critic = tch::no_grad(|| model(&s_states.get(-1)).0);
             r.get(-1).copy_(&critic.view([NPROCS]));
             for s in (0..NSTEPS).rev() {
@@ -125,8 +125,7 @@ pub fn train() -> cpython::PyResult<()> {
             let index = s_actions.unsqueeze(-1).to_device(device);
             log_probs.gather(2, &index, false).squeeze_dim(-1)
         };
-        let dist_entropy =
-            (-log_probs * probs).sum_dim_intlist(Some([-1].as_slice()), false, Float).mean(Float);
+        let dist_entropy = (-log_probs * probs).sum_dim_intlist(-1, false, Float).mean(Float);
         let advantages = s_returns.narrow(0, 0, NSTEPS).to_device(device) - critic;
         let value_loss = (&advantages * &advantages).mean(Float);
         let action_loss = (-advantages.detach() * action_log_probs).mean(Float);

examples/reinforcement-learning/ddpg.rs

@@ -48,13 +48,13 @@ struct OuNoise {
 
 impl OuNoise {
     fn new(mu: f64, theta: f64, sigma: f64, num_actions: usize) -> Self {
-        let state = Tensor::ones(&[num_actions as _], FLOAT_CPU);
+        let state = Tensor::ones([num_actions as _], FLOAT_CPU);
         Self { mu, theta, sigma, state }
     }
 
     fn sample(&mut self) -> &Tensor {
         let dx = self.theta * (self.mu - &self.state)
-            + self.sigma * Tensor::randn(&self.state.size(), FLOAT_CPU);
+            + self.sigma * Tensor::randn(self.state.size(), FLOAT_CPU);
         self.state += dx;
         &self.state
     }
@@ -73,10 +73,10 @@ struct ReplayBuffer {
 impl ReplayBuffer {
     fn new(capacity: usize, num_obs: usize, num_actions: usize) -> Self {
         Self {
-            obs: Tensor::zeros(&[capacity as _, num_obs as _], FLOAT_CPU),
-            next_obs: Tensor::zeros(&[capacity as _, num_obs as _], FLOAT_CPU),
-            rewards: Tensor::zeros(&[capacity as _, 1], FLOAT_CPU),
-            actions: Tensor::zeros(&[capacity as _, num_actions as _], FLOAT_CPU),
+            obs: Tensor::zeros([capacity as _, num_obs as _], FLOAT_CPU),
+            next_obs: Tensor::zeros([capacity as _, num_obs as _], FLOAT_CPU),
+            rewards: Tensor::zeros([capacity as _, 1], FLOAT_CPU),
+            actions: Tensor::zeros([capacity as _, num_actions as _], FLOAT_CPU),
             capacity,
             len: 0,
             i: 0,
@@ -101,7 +101,7 @@ impl ReplayBuffer {
         }
 
         let batch_size = batch_size.min(self.len - 1);
-        let batch_indexes = Tensor::randint((self.len - 2) as _, &[batch_size as _], INT64_CPU);
+        let batch_indexes = Tensor::randint((self.len - 2) as _, [batch_size as _], INT64_CPU);
 
         let states = self.obs.index_select(0, &batch_indexes);
         let next_states = self.next_obs.index_select(0, &batch_indexes);

examples/reinforcement-learning/policy_gradient.rs

@@ -72,14 +72,11 @@ pub fn run() -> cpython::PyResult<()> {
         let rewards = accumulate_rewards(&steps);
         let rewards = Tensor::of_slice(&rewards).to_kind(Float);
         let action_mask =
-            Tensor::zeros(&[batch_size, 2], tch::kind::FLOAT_CPU).scatter_value(1, &actions, 1.0);
+            Tensor::zeros([batch_size, 2], tch::kind::FLOAT_CPU).scatter_value(1, &actions, 1.0);
         let obs: Vec<Tensor> = steps.into_iter().map(|s| s.obs).collect();
         let logits = Tensor::stack(&obs, 0).apply(&model);
-        let log_probs = (action_mask * logits.log_softmax(1, Float)).sum_dim_intlist(
-            Some([1].as_slice()),
-            false,
-            Float,
-        );
+        let log_probs =
+            (action_mask * logits.log_softmax(1, Float)).sum_dim_intlist(1, false, Float);
         let loss = -(rewards * log_probs).mean(Float);
         opt.backward_step(&loss)
     }

examples/reinforcement-learning/ppo.rs

@@ -49,7 +49,7 @@ struct FrameStack {
 
 impl FrameStack {
     fn new(nprocs: i64, nstack: i64) -> FrameStack {
-        FrameStack { data: Tensor::zeros(&[nprocs, nstack, 84, 84], FLOAT_CPU), nprocs, nstack }
+        FrameStack { data: Tensor::zeros([nprocs, nstack, 84, 84], FLOAT_CPU), nprocs, nstack }
     }
 
     fn update<'a>(&'a mut self, img: &Tensor, masks: Option<&Tensor>) -> &'a Tensor {
@@ -75,21 +75,21 @@ pub fn train() -> cpython::PyResult<()> {
     let model = model(&vs.root(), env.action_space());
     let mut opt = nn::Adam::default().build(&vs, 1e-4).unwrap();
 
-    let mut sum_rewards = Tensor::zeros(&[NPROCS], FLOAT_CPU);
+    let mut sum_rewards = Tensor::zeros([NPROCS], FLOAT_CPU);
     let mut total_rewards = 0f64;
     let mut total_episodes = 0f64;
 
     let mut frame_stack = FrameStack::new(NPROCS, NSTACK);
     let _ = frame_stack.update(&env.reset()?, None);
-    let s_states = Tensor::zeros(&[NSTEPS + 1, NPROCS, NSTACK, 84, 84], FLOAT_CPU);
+    let s_states = Tensor::zeros([NSTEPS + 1, NPROCS, NSTACK, 84, 84], FLOAT_CPU);
 
     let train_size = NSTEPS * NPROCS;
     for update_index in 0..UPDATES {
         s_states.get(0).copy_(&s_states.get(-1));
-        let s_values = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
-        let s_rewards = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
-        let s_actions = Tensor::zeros(&[NSTEPS, NPROCS], INT64_CPU);
-        let s_masks = Tensor::zeros(&[NSTEPS, NPROCS], FLOAT_CPU);
+        let s_values = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
+        let s_rewards = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
+        let s_actions = Tensor::zeros([NSTEPS, NPROCS], INT64_CPU);
+        let s_masks = Tensor::zeros([NSTEPS, NPROCS], FLOAT_CPU);
         for s in 0..NSTEPS {
             let (critic, actor) = tch::no_grad(|| model(&s_states.get(s)));
             let probs = actor.softmax(-1, Kind::Float);
@@ -111,7 +111,7 @@ pub fn train() -> cpython::PyResult<()> {
         }
         let states = s_states.narrow(0, 0, NSTEPS).view([train_size, NSTACK, 84, 84]);
         let returns = {
-            let r = Tensor::zeros(&[NSTEPS + 1, NPROCS], FLOAT_CPU);
+            let r = Tensor::zeros([NSTEPS + 1, NPROCS], FLOAT_CPU);
             let critic = tch::no_grad(|| model(&s_states.get(-1)).0);
             r.get(-1).copy_(&critic.view([NPROCS]));
             for s in (0..NSTEPS).rev() {
@@ -122,7 +122,7 @@ pub fn train() -> cpython::PyResult<()> {
         };
         let actions = s_actions.view([train_size]);
         for _index in 0..OPTIM_EPOCHS {
-            let batch_indexes = Tensor::randint(train_size, &[OPTIM_BATCHSIZE], INT64_CPU);
+            let batch_indexes = Tensor::randint(train_size, [OPTIM_BATCHSIZE], INT64_CPU);
             let states = states.index_select(0, &batch_indexes);
             let actions = actions.index_select(0, &batch_indexes);
             let returns = returns.index_select(0, &batch_indexes);
@@ -133,9 +133,8 @@ pub fn train() -> cpython::PyResult<()> {
                 let index = actions.unsqueeze(-1).to_device(device);
                 log_probs.gather(-1, &index, false).squeeze_dim(-1)
             };
-            let dist_entropy = (-log_probs * probs)
-                .sum_dim_intlist(Some([-1].as_slice()), false, Kind::Float)
-                .mean(Kind::Float);
+            let dist_entropy =
+                (-log_probs * probs).sum_dim_intlist(-1, false, Kind::Float).mean(Kind::Float);
             let advantages = returns.to_device(device) - critic;
             let value_loss = (&advantages * &advantages).mean(Kind::Float);
             let action_loss = (-advantages.detach() * action_log_probs).mean(Kind::Float);