get_probabilistic_reparameterization_input_transform gives incorrect dimension when optimising over only categorical inputs

[exs-hkenlay]

I'm trying to optimise a function which just has categorical inputs. However I think the get_probabilistic_reparameterization_input_transform functions appears to fail in this case.

Here is an example with one continuous feature and two categorical features.

import torch
from discrete_mixed_bo.probabilistic_reparameterization import get_probabilistic_reparameterization_input_transform

input_transform = get_probabilistic_reparameterization_input_transform(
    dim=13, ##
    use_analytic=True,
    integer_indices=[],
    integer_bounds=torch.zeros((2, 0), dtype=torch.float64),
    categorical_features=OrderedDict([(1, 6), (2, 6)]),
    tau=0.1,
)

print(input_transform(torch.rand(100, 1, 1, 13)).shape)

This will output torch.Size([100, 36, 1, 13]). Now if we remove the continuous feature.

import torch
from discrete_mixed_bo.probabilistic_reparameterization import get_probabilistic_reparameterization_input_transform

input_transform = get_probabilistic_reparameterization_input_transform(
    dim=12, ##
    use_analytic=True,
    integer_indices=[],
    integer_bounds=torch.zeros((2, 0), dtype=torch.float64),
    categorical_features=OrderedDict([(0, 6), (1, 6)]),
    tau=0.1,
)

print(input_transform(torch.rand(100, 1, 1, 12)).shape)

And the output is torch.Size([36, 12]). Should this not be torch.Size([100, 36, 1, 12])?

I've managed to find the line that seems to cause this issue. It is the tf.eval() on line 95. I'm not sure why this causes the transform to collapse the dimensions.

[sdaulton]

Hi @hkenlay, Thanks for raising the question. I didn't consider categorical-only search spaces in the original implementation in this research repo. I am working on upstreaming probabilistic reparameterization into botorch and the pull request upstream (botorch/pull/1533) does support categorical-only search spaces. I have attached a notebook demonstrating this functionality using that pull request. Let me know if you run into any issues (as the pull request is in development). I intend to use the upstream version in this repo, once the pull requests land. Thanks!
discrete_mixed_bo_categorical.ipynb.txt

[exs-hkenlay]

Hi @sdaulton, thank you for your reply. Unfortunately the notebook has a few issues for me. I can describe them here or in the PR (in which case can you upload the notebook there for context?)?

[sdaulton]

Let's chat about it on the pull request in botorch. I will upload the nb there

[exs-hkenlay]

Thanks @sdaulton.