Variable Dimensions

Since SOL v0.5 we support variable dimensions (VDims). When you run sol.optimize, SOL will automatically detect variable dimensions. The acceptable input shapes get printed by SOL after analyzing the model.

Inputs:  in_0 [#0, 5, #1, #2, #3]
Outputs: out_0  {
    "A": [#0, 5, #1, #2, #3],
    "B": [#0, 5, #1, 3, 3],
    "C": [#0, 5, #1, 5, 7],

Here we see that we have 4 VDims (#0 to #3). Depending on the structure of your neural network, SOL will restrict certain dimensions to be of fixed size. For performance reasons, SOL disables all VDims, so you need to enable them by hand using:

sol_model = sol.optimize(model, [torch.rand(5, 5, 5, 5, 5)], vdims=[True, False, 1])

This enables the #0 to accept any size. #1 will only accept the size that was used when parsing the model (the default behavior). #2 will only accept the size to be 1. As #3 does not get set, it will use the default behavior. So the compatible shape for this example is [*, 5, 5, 1, 5].

Setting the VDims needs to happen BEFORE you execute the model for the first time, otherwise SOL will not obey your settings! If you need to change your VDims call sol.cache.clear() at the beginning of your script.