PERF: change repeat to expand in chunk processing

mdw771 · mdw771 · commit 176c3e64007b · 2025-10-08T14:34:51.000-05:00
diff --git a/src/ptychi/utils.py b/src/ptychi/utils.py
@@ -412,7 +412,7 @@ def chunked_processing(
         current_chunk_size = kwargs_chunk[list(kwargs_chunk.keys())[0]].shape[0]
         if replicated_kwargs is not None:
             replicated_kwargs_chunk = {
-                key: torch.repeat_interleave(value, current_chunk_size, dim=0)
+                key: value.expand(current_chunk_size, *value.shape[1:])
                 for key, value in replicated_kwargs.items()
             }
             kwargs_chunk.update(replicated_kwargs_chunk)

Original file line number	Diff line number	Diff line change
`@@ -412,7 +412,7 @@ def chunked_processing(`
`412`	`412`	`current_chunk_size = kwargs_chunk[list(kwargs_chunk.keys())[0]].shape[0]`
`413`	`413`	`if replicated_kwargs is not None:`
`414`	`414`	`replicated_kwargs_chunk = {`
`415`		`- key: torch.repeat_interleave(value, current_chunk_size, dim=0)`
	`415`	`+ key: value.expand(current_chunk_size, *value.shape[1:])`
`416`	`416`	`for key, value in replicated_kwargs.items()`
`417`	`417`	`}`
`418`	`418`	`kwargs_chunk.update(replicated_kwargs_chunk)`