[Bug]: StratifiedStandardize get_task_value_remapping fails when attempting to re-map tasks

[Hrovatin]

What happened?

If tasks passed to StratifiedStandardize are a range(N_tasks) then the StratifiedStandardize does not need to -re-map them to the range and the scaling succeeds:

# Works
botorch.models.transforms.outcome.StratifiedStandardize(
         task_values=torch.tensor([0,1]).to(torch.long),
         stratification_idx=0,
)

However, if they do not match the range, e.g. are some other values, the StratifiedStandardize attempts to re-map them to a range, which raises an error when creating the re-mapping map.

botorch.models.transforms.outcome.StratifiedStandardize(
         task_values=torch.tensor([1,2]).to(torch.long),
         stratification_idx=0,
)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[43], line 1
----> 1 botorch.models.transforms.outcome.StratifiedStandardize(
      2          task_values=torch.tensor([1,2]).to(torch.long),
      3          stratification_idx=0,
      4 )

File [~/miniforge3/envs/baybe/lib/python3.10/site-packages/botorch/models/transforms/outcome.py:535](http://localhost:8888/lab/workspaces/auto-J/tree/baybe-analysis/method_improvements/~/miniforge3/envs/baybe/lib/python3.10/site-packages/botorch/models/transforms/outcome.py#line=534), in StratifiedStandardize.__init__(self, task_values, stratification_idx, batch_shape, min_stdv)
    533 self._stratification_idx = stratification_idx
    534 task_values = task_values.unique(sorted=True)
--> 535 self.strata_mapping = get_task_value_remapping(task_values, dtype=torch.long)
    536 if self.strata_mapping is None:
    537     self.strata_mapping = task_values

File [~/miniforge3/envs/baybe/lib/python3.10/site-packages/botorch/models/utils/assorted.py:426](http://localhost:8888/lab/workspaces/auto-J/tree/baybe-analysis/method_improvements/~/miniforge3/envs/baybe/lib/python3.10/site-packages/botorch/models/utils/assorted.py#line=425), in get_task_value_remapping(task_values, dtype)
    422 mapper = None
    423 if not torch.equal(task_values, task_range):
    424     # Create a tensor that maps task values to new task values.
    425     # The number of tasks should be small, so this should be quite efficient.
--> 426     mapper = torch.full(
    427         (int(task_values.max().item()) + 1,),
    428         float("nan"),
    429         dtype=dtype,
    430         device=task_values.device,
    431     )
    432     mapper[task_values] = task_range.to(dtype=dtype)
    433 return mapper

RuntimeError: value cannot be converted to type int64_t without overflow

This occurs because the utils code attempts to create long-typed tensor that is filled with nan values. Long-typed tensors don't seem to be able to contain nan values.

torch.full(
         (2,),
      float("nan"),
      dtype=torch.long,
        device=task_values.device,
 )
RuntimeError: value cannot be converted to type int64_t without overflow

# Works if dtype is not int/long
torch.full(
         (2,), # Number of tasks passed to thefunction
          float("nan"),
        device=task_values.device,
 )
tensor([nan, nan])

The issue occurs here: https://github.com/pytorch/botorch/blob/9a7c5176e14de1a18b798c7291cc9fbbcbc96d5b/botorch/models/utils/assorted.py#L432C18-L437C10

Please provide a minimal, reproducible example of the unexpected behavior.

See above

Please paste any relevant traceback/logs produced by the example provided.

BoTorch Version

0.13.0

Python Version

No response

Operating System

No response

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[esantorella]

Thanks for the clear repro and explanation! We'll look into this.

[sdaulton]

Thanks for raising this. I'll put up a PR to fix this.

[sdaulton]

#2757