refactor: Avoid cat in Jac when not needed (#518)

Author: ValerianRey

CHANGELOG.md

@@ -35,6 +35,8 @@ changelog does not include internal changes that do not affect the user.
   jac_to_grad(shared_module.parameters(), aggregator)
   ```
 
+- Removed an unnecessary memory duplication. This should significantly improve the memory efficiency
+  of `autojac`.
 - Removed an unnecessary internal cloning of gradient. This should slightly improve the memory
   efficiency of `autojac`.
 

src/torchjd/autojac/_transform/_jac.py

@@ -91,7 +91,12 @@ def _differentiate(self, jac_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
         jacs_chunks.append(_get_jacs_chunk(jac_outputs_chunk, get_vjp_last))
 
         n_inputs = len(self.inputs)
-        jacs = tuple(torch.cat([chunks[i] for chunks in jacs_chunks]) for i in range(n_inputs))
+        if len(jacs_chunks) == 1:
+            # Avoid using cat to avoid doubling memory usage, if it's not needed
+            jacs = jacs_chunks[0]
+        else:
+            jacs = tuple(torch.cat([chunks[i] for chunks in jacs_chunks]) for i in range(n_inputs))
+
         return jacs