Use int64 when calculating data size in split_acquisition_function #795
Conversation
| @@ -48,7 +48,8 @@ def wrapper(x: TensorType) -> TensorType: | |||
| if length == 0: | |||
| return fn(x) | |||
|
|
|||
| elements_per_block = tf.size(x) / length | |||
| # Use int64 to calculate the input tensor size, otherwise we can overflow for large tensors. | |||
There was a problem hiding this comment.
this is good but could you please add in the docstrings few sentences about memory requirements wrt scaling?
There was a problem hiding this comment.
Do you mean docstring for split_acquisition_function? Could you give an example please? I could say that the memory usage of tensor x is its flattened size times 8, but that is fairly obvious.
Or do you mean in the optimizer docstring?
There was a problem hiding this comment.
i think it would be more appropriate in the generate_continuous_optimizer docstring where tiled_candidates becomes very big - I was thinking of an additional note with a warning that using a large num_initial_samples with high dimensional search space and/or batch sizes will require a lot of memory.
There was a problem hiding this comment.
I have added a note.
Related issue(s)/PRs: None
Summary
Use
int64to calculate the input tensor size when splitting acquisition functions. Otherwise, we overflow for large tensors, e.g. of shape[30000, 1000, 100].It is tricky to add a unit test to cover the limit, as we run into memory issues. Using mocking or sparse-tensors gets a bit messy, so skipping unit testing for this simple change.
Fully backwards compatible: yes
PR checklist