Using botorch.acquisition.multi_step_lookahead : RuntimeError: split_with_sizes expects split_sizes have only non-negative entries

[Adalia456]

Hi everyone,

I'm currently working on integrating the qMultiStepLookahead acquisition function into my project using BoTorch, following a tutorial as a beginner in this domain. This code snippet is part of my initial exploration and testing phase before full implementation.

I've encountered an error that I'm struggling to resolve. I suspect it may be due to incorrect parameters passed to the acquisition function. However, I'm having difficulty pinpointing the exact issue.

I would greatly appreciate any insights or suggestions on how to debug and correct this issue.

Thank you for your help!

Here's the error message I'm encountering, and the part of my code concerning my acquisition function :

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[10], line 59
     50 """qmc_sampler = SobolQMCNormalSampler(sample_shape=torch.Size([MC_SAMPLES]))
     51 qEI = qExpectedImprovement(
     52         model=model,
   (...)
     55         best_f = best_observed_value_nei
     56     )"""
     58 # Optimize the acquisition function and get new candidate observations, manually entered by the user 
---> 59 new_x_nei = optimize_acqf_and_get_observation(qMultiStepsLookahead)
     60 print(new_x_nei)
     61 new_obj_nei1 = torch.tensor([[float(x) for x in input("Enter 1 float values separated by spaces: ").split()]])

Cell In[8], line 24, in optimize_acqf_and_get_observation(acq_func)
     14 """
     15 Optimizes the acquisition function and returns a new candidate observation.
     16 
   (...)
     21     new_x: Tensor containing the new candidate observation.
     22 """
     23 # Optimize the acquisition function
---> 24 candidates, _ = optimize_acqf(
     25     acq_function=acq_func,
     26     bounds=bounds,
     27     q=BATCH_SIZE,
     28     num_restarts=NUM_RESTARTS,
     29     raw_samples=RAW_SAMPLES,  # Used for initialization heuristic
     30     options={"batch_limit": 1, "maxiter": 200},
     31 )
     32 # Get the new candidate observation
     33 new_x = candidates.detach()

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:567, in optimize_acqf(acq_function, bounds, q, num_restarts, raw_samples, options, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, fixed_features, post_processing_func, batch_initial_conditions, return_best_only, gen_candidates, sequential, ic_generator, timeout_sec, return_full_tree, retry_on_optimization_warning, **ic_gen_kwargs)
    544     gen_candidates = gen_candidates_scipy
    545 opt_acqf_inputs = OptimizeAcqfInputs(
    546     acq_function=acq_function,
    547     bounds=bounds,
   (...)
    565     ic_gen_kwargs=ic_gen_kwargs,
    566 )
--> 567 return _optimize_acqf(opt_acqf_inputs)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:588, in _optimize_acqf(opt_inputs)
    585     return _optimize_acqf_sequential_q(opt_inputs=opt_inputs)
    587 # Batch optimization (including the case q=1)
--> 588 return _optimize_acqf_batch(opt_inputs=opt_inputs)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:275, in _optimize_acqf_batch(opt_inputs)
    272     batch_initial_conditions = opt_inputs.batch_initial_conditions
    273 else:
    274     # pyre-ignore[28]: Unexpected keyword argument `acq_function` to anonymous call.
--> 275     batch_initial_conditions = opt_inputs.get_ic_generator()(
    276         acq_function=opt_inputs.acq_function,
    277         bounds=opt_inputs.bounds,
    278         q=opt_inputs.q,
    279         num_restarts=opt_inputs.num_restarts,
    280         raw_samples=opt_inputs.raw_samples,
    281         fixed_features=opt_inputs.fixed_features,
    282         options=options,
    283         inequality_constraints=opt_inputs.inequality_constraints,
    284         equality_constraints=opt_inputs.equality_constraints,
    285         **opt_inputs.ic_gen_kwargs,
    286     )
    288 batch_limit: int = options.get(
    289     "batch_limit",
    290     (
   (...)
    294     ),
    295 )
    297 def _optimize_batch_candidates() -> Tuple[Tensor, Tensor, List[Warning]]:

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\initializers.py:417, in gen_batch_initial_conditions(acq_function, bounds, q, num_restarts, raw_samples, fixed_features, options, inequality_constraints, equality_constraints, generator, fixed_X_fantasies)
    415 while start_idx < X_rnd.shape[0]:
    416     end_idx = min(start_idx + batch_limit, X_rnd.shape[0])
--> 417     Y_rnd_curr = acq_function(
    418         X_rnd[start_idx:end_idx].to(device=device)
    419     ).cpu()
    420     Y_rnd_list.append(Y_rnd_curr)
    421     start_idx += batch_limit

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\nn\modules\module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\nn\modules\module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\utils\transforms.py:259, in t_batch_mode_transform.<locals>.decorator.<locals>.decorated(acqf, X, *args, **kwargs)
    257 # add t-batch dim
    258 X = X if X.dim() > 2 else X.unsqueeze(0)
--> 259 output = method(acqf, X, *args, **kwargs)
    260 if hasattr(acqf, "model") and is_ensemble(acqf.model):
    261     # IDEA: this could be wrapped into SampleReducingMCAcquisitionFunction
    262     output = (
    263         output.mean(dim=-1) if not acqf._log else logmeanexp(output, dim=-1)
    264     )

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\acquisition\multi_step_lookahead.py:179, in qMultiStepLookahead.forward(self, X)
    166 @t_batch_mode_transform()
    167 def forward(self, X: Tensor) -> Tensor:
    168     r"""Evaluate qMultiStepLookahead on the candidate set X.
    169 
    170     Args:
   (...)
    177         The acquisition value for each batch as a tensor of shape `batch_shape`.
    178     """
--> 179     Xs = self.get_multi_step_tree_input_representation(X)
    181     # set batch_range on samplers if not collapsing on fantasy dims
    182     if not self._collapse_fantasy_base_samples:

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\acquisition\multi_step_lookahead.py:269, in qMultiStepLookahead.get_multi_step_tree_input_representation(self, X)
    266 batch_shape, shapes, sizes = self.get_split_shapes(X=X)
    267 # Each X_i in Xsplit has shape batch_shape x qtilde x d with
    268 # qtilde = f_i * ... * f_1 * q_i
--> 269 Xsplit = torch.split(X, sizes, dim=-2)
    270 # now reshape (need to permute batch_shape and qtilde dimensions for i > 0)
    271 perm = [-2] + list(range(len(batch_shape))) + [-1]

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\functional.py:195, in split(tensor, split_size_or_sections, dim)
    189     return handle_torch_function(
    190         split, (tensor,), tensor, split_size_or_sections, dim=dim)
    191 # Overwriting reason:
    192 # This dispatches to two ATen functions depending on the type of
    193 # split_size_or_sections. The branching code is in _tensor.py, which we
    194 # call here.
--> 195 return tensor.split(split_size_or_sections, dim)

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\_tensor.py:921, in Tensor.split(self, split_size, dim)
    919     return torch._VF.split(self, split_size, dim)  # type: ignore[attr-defined]
    920 else:
--> 921     return torch._VF.split_with_sizes(self, split_size, dim)

RuntimeError: split_with_sizes expects split_sizes have only non-negative entries, but got split_sizes=[-[1000999, 1000, 1000000](callto:1000999, 1000, 1000000)]```

> Here is how I coded my botorch loop : 

```import time
import warnings
import torch
from botorch import fit_gpytorch_mll
from botorch.acquisition.multi_step_lookahead import qMultiStepLookahead
from botorch.acquisition import qExpectedImprovement
from botorch.exceptions import BadInitialCandidatesWarning
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.acquisition import PosteriorMean

# Suppress specific warnings
warnings.filterwarnings("ignore", category=BadInitialCandidatesWarning)
warnings.filterwarnings("ignore", category=RuntimeWarning)

# Define the number of trials, batches, and Monte Carlo samples based on the SMOKE_TEST flag
N_TRIALS = 3 if not SMOKE_TEST else 2
N_BATCH = 20 if not SMOKE_TEST else 2
MC_SAMPLES = 256 if not SMOKE_TEST else 32

verbose = False

# List to store the best observed values for all trials
best_observed_nei = []

# Generate initial training data and initialize the model
train_x, train_obj_nei, best_observed_value_nei = generate_initial_data(n=10)
mll_nei, model = initialize_model(train_x, train_obj_nei)

# Append the best observed initial value to the list
best_observed_nei.append(best_observed_value_nei)

# Run N_BATCH rounds of Bayesian Optimization after the initial random batch
for iteration in range(1, N_BATCH + 1):

    # Fit the model to the current data
    fit_gpytorch_mll(mll_nei)

    # Create the qMultiStepLookahead object
    qMultiStepsLookahead = qMultiStepLookahead(
        model=model,
        batch_sizes=[1,1],
        num_fantasies = [1000,1000],
        valfunc_cls=[None,None,PosteriorMean],   # Using PosteriorMean as default
        valfunc_argfacs=[None,None,None],  # Pass None as default
        objective=None,
        posterior_transform=None,
        X_pending=None,
        collapse_fantasy_base_samples=True
    )

    # Optimize the acquisition function and get new candidate observations, manually entered by the user
    new_x_nei = optimize_acqf_and_get_observation(qMultiStepsLookahead)
    print(new_x_nei)
    new_obj_nei1 = torch.tensor([[float(x) for x in input("Enter 1 float values separated by spaces: ").split()]])
    """new_obj_nei2 = torch.tensor([[float(x) for x in input("Enter 1 float values separated by spaces: ").split()]])"""
    # Update training points with the new observations
    
    train_x = torch.cat([train_x, new_x_nei])
    train_obj_nei = torch.cat([train_obj_nei, new_obj_nei1])
    
    # Update the best observed value
    best_value_nei = train_x.max().item()
    best_observed_nei.append(best_value_nei)

    # Reinitialize the model with the current state dict to prepare for the next iteration
    mll, model = initialize_model(
        train_x,
        train_obj_nei,
    )```

[esantorella]

Could you provide a reproducible example, including the definitions of SMOKE_TEST, generate_initial_data, and anything else that would be needed to run the code?