Does Ax allow attachment of trials where trails have input values outside the experiment's parameter bounds?

[jakobzeitler]

A possible Ax data loading routine might look like this:

trial_data_Y = train_Y.numpy()
for i, x in enumerate(train_X.numpy()):
    trial_data = {}
    for j, input in enumerate(inputs):
      trial_data[input] = value
    trial_complete_data = {}
    for j, output in enumerate(outputs):
      trial_complete_data[output] = trial_data_Y[i,j]
    p, trial_index = ax_client.attach_trial(parameters=trial_data)
    ax_client.complete_trial(trial_index=trial_index, raw_data=trial_complete_data)

If the bounds set for the experiment in the Ax client are tighter than sample values in the some of the trials loaded via the routine above, an error will be raised by this validation step in search_space.py

Is this intended behaviour?

I fearing I might be abusing Ax API here, but at least theoretically, we should able to sample from posteriors generated from models with data where constraints are tighter than samples used to train the surrogate.

Thanks a lot in advance!

[Balandat]

we should able to sample from posteriors generated from models with data where constraints are tighter than samples used to train the surrogate.

Yes, we definitely want to support this. #3040 is relevant, which separates concepts of "modeling space" and "search space". From that PR:

By default, modeling space will be expanded to cover the full range of RangeParameter values seen in the training data. This means that training data will never be marked OOD for falling outside RangeParameter box bounds. [...] At gen time, we will still respect the search space bounds as given.

So we do support this internally, but it appears the higher level API has some (potentially unnecessary; could warn instead) validation step. Could you share the full stack trace of the error? I'm having a bit of a hard time manually tracking down how this error is actually raised from your code entry point.

[jakobzeitler]

Hi Max,

Very glad to see I am not going crazy.

I have put together this Google colab for you which you can clone and play with, follow traces etc. :
https://colab.research.google.com/drive/1PufYdMZlZ3oj6qolinVGyqhlA3AGzsoJ?usp=sharing

The following lines will run:

trial_data_X = {"x1":0.0}
trial_data_y = {"hartmann6":0.0}

p, trial_index = ax_client.attach_trial(parameters=trial_data_X)
ax_client.complete_trial(trial_index=trial_index, raw_data=trial_data_y)

While these lines will not, due to the search_space validator, because "2.0" is outside the bounds of [0.0,1.0]:

trial_data_X = {"x1":2.0}
trial_data_y = {"hartmann6":0.0}

p, trial_index = ax_client.attach_trial(parameters=trial_data_X)
ax_client.complete_trial(trial_index=trial_index, raw_data=trial_data_y)

We'd like to get this sorted ASAP, shall I continue where @bletham has left off (had a quick look, so not sure how far he got), make changes and submit a merge-request? Or should we aim for some more coordination?

Thanks a lot in advance!

[Balandat]

shall I continue where bletham has left off

This addressed the internals of this on the modeling side; we essentially just need to change the user-facing API here.

We're working on a new, consolidated API (preview here: https://github.com/facebook/Ax/blob/main/ax/preview/api/client.py) and the right thing to do would be to ensure things are handled properly there - cc @mpolson64

It'll be a bit until this is fully released, so one could consider updating the existing API to allow for this. Let me discuss this with some folks this week.

[jakobzeitler]

Thanks for the update and context, we'll attempt a quick fix on the current API, and looking forward to what you conclude with your folks. :)

[jakobzeitler]

We are starting on fixing the current API this week, any news from your side? @Balandat

[Balandat]

@mpolson64 was going to look into this

[lena-kashtelyan]

Hi @jakobzeitler! This is currently high priority in our backlog, but we unfortunately cannot attend to this immediately.

[josephinecowley]

Hi! I've implemented changes to expose the raise_error bool in search_space.validate_membership() to the client using a new parameter raise_parameter_error in ax_client.attach_trial(). This made the most sense to me as it requires minimal changes and shouldn't break existing code. This value is default to True (same behaviour as before), so the client will have to set it to False if they want to allow the behaviour described by @jakobzeitler.

[mpolson64]

Hi @josephinecowley thanks for putting this together! This seems like a fairly safe change an should be alright for us to land, but let me see what I can do to get this larger solution @Balandat @lena-kashtelyan and I have been thinking about prioritized and implemented, ideally within the week.

If it looks like it is going to get delayed longer than that we can make sure to get #3402 or something to its effect in instead.