Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| from typing import Any, Callable | |
| def run_multiseed(fn: Callable, **kwargs: list[Any] | Any) -> Any: | |
| """Run a callable once per seed, indexing all list kwargs in lockstep. | |
| All keyword arguments must be **either** all lists or all non-lists. | |
| Mixing is not allowed. | |
| - **All non-list**: ``fn`` is called once with the kwargs as-is, and its | |
| return value is passed through directly. | |
| - **All list**: every list must have the same length *N*. ``fn`` is called | |
| *N* times — call *i* receives ``{k: v[i] for k, v in kwargs}``. | |
| Results are collected into a list. If ``fn`` returns a tuple, the | |
| results are transposed into a tuple of lists. | |
| Args: | |
| fn: Callable to invoke per seed. | |
| **kwargs: Keyword arguments for ``fn``. Must be **all lists** (one | |
| element per seed, all the same length) or **all non-lists** (single | |
| call). | |
| Returns: | |
| - All non-list kwargs: the raw return value of ``fn``. | |
| - All list kwargs, ``fn`` returns a tuple: a ``tuple`` of lists, | |
| transposed across calls. | |
| - All list kwargs, ``fn`` returns non-tuple: a ``list`` of return | |
| values. | |
| Raises: | |
| AssertionError: If kwargs mix lists and non-lists, or if list kwargs | |
| have differing lengths. | |
| Examples: | |
| Single call (no lists):: | |
| run_multiseed(lambda x, y: x + y, x=1, y=2) # returns 3 | |
| Multiple calls with all-list kwargs:: | |
| run_multiseed(lambda x, y: x * y, x=[1, 2, 3], y=[10, 20, 30]) | |
| # returns [10, 40, 90] | |
| Tuple return transposition:: | |
| run_multiseed(lambda x: (x, -x), x=[1, 2]) | |
| # returns ([1, 2], [-1, -2]) | |
| """ | |
| all_list = all(isinstance(v, list) for v in kwargs.values()) | |
| all_non_list = all(not isinstance(v, list) for v in kwargs.values()) | |
| assert all_list or all_non_list, "All kwargs must be lists or all must be non-lists, cannot mix" | |
| if all_non_list: | |
| return fn(**kwargs) | |
| lengths = {len(v) for v in kwargs.values()} | |
| assert len(lengths) == 1, f"All list arguments must have the same length, got {lengths}" | |
| num_calls = lengths.pop() | |
| results = [] | |
| for i in range(num_calls): | |
| kwargs_i = {k: v[i] for k, v in kwargs.items()} | |
| results.append(fn(**kwargs_i)) | |
| if results and isinstance(results[0], tuple): | |
| return tuple(list(items) for items in zip(*results)) | |
| return results | |