feat: Add batched tensor expansion utilities for O(S) ensemble unroll…

[Debadri-das]

Describe your changes

This PR introduces a standalone batch-expanding PyTorch utility (expand_ensemble_batch and fold_ensemble_batch) to optimize tensor manipulation for Ensemble Lateral Boundary Conditions (LBCs) and probabilistic forecasting variables.

Motivation and Context: Currently, probabilistic processing of $S$ ensemble members often relies on sequential for loops. To support the upcoming model refactoring and probabilistic capabilities without introducing $O(S)$ sequential bottlenecks, these functions leverage torch.repeat_interleave() and .view() to dynamically flatten state tensors (B, ...) -> (B * S, ...). This enables optimized parallelized operations over the GPU without permanently mutating standard AR model dimensions.

Dependencies: None. Built strictly with standard PyTorch tensor primitives.

Issue Link

#49 and #62

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch .
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug
  • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

[Sir-Sloth-The-Lazy]

Just a heads up you missed the case when S == T . Love the idea !

[Sir-Sloth-The-Lazy]

is this intentional ?

[sadamov]

linking #226 and #332 here, to remember the potential connection. this is one for @joeloskarsson to review.

[Debadri-das]

is this intentional ?
No, the code was not intentional at all, it was a flaw in the design, which I fixed now.

[Sir-Sloth-The-Lazy]

@Debadri-das would you be working on this further ? It would be really great if you do or I can take over if you have other work going on. I have this as a dependency for my work :) Hope to hear from you soon !

[joeloskarsson]

Thanks, added a couple thoughts. Do also check out @Sir-Sloth-The-Lazy's comment, as he is eager to make use of these methods.

[joeloskarsson]

It seems like it would be useful for this to be able to take either the number of members or the batch size (and if both are given, check that their product is size of first dim).

[joeloskarsson]

I think it would be better to have two separate methods for this, as these are two very different things (repeating the same state vs just a reshape). If there is a lot of repeated code add a supporting helper function to avoid this.

[joeloskarsson]

Happy to discuss if you see a good argument for this being one function!

[joeloskarsson]

The added methods here do something very close to expand_to_batch

neural-lam/neural_lam/models/step_predictors/base.py

Lines 86 to 103 in a208ada

	def expand_to_batch(self, x: torch.Tensor, batch_size: int) -> torch.Tensor:
	"""
	Expand a shared node-feature tensor into a batch of copies.
	Parameters
	----------
	x : torch.Tensor
	Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number
	of nodes and ``d`` is the feature dimension.
	batch_size : int
	Target batch size ``B``.
	Returns
	-------
	torch.Tensor
	Shape ``(B, N, d)``. Batch-expanded view of ``x``.
	"""
	return x.unsqueeze(0).expand(batch_size, -1, -1)

and it seems suboptimal that these methods are in different places in the codebase. Actually, expand_to_batch seems like it could just as well sit in utils.py, and there is no reason why it has to be a class method. I think it would make sense to in this PR also move that to utils, to have all batch-dim reordering/reshaping functions in the same place.

[sadamov]

Closing this in favour of #649, which consolidates the probabilistic metrics + ensemble plumbing track for v0.8.0. The expand_ensemble_batch / fold_ensemble_batch utilities will be folded into the consolidated PR. Thanks for the contribution @Debadri-das, you'll be credited.