SplitToSequnce const folding only works for static shapes

[yuanyao-nv]

In diffusers v0.35, the FluxTransformer model definition introduced split_with_sizes (here)

When dynamic shape is turned on the onnxscript const folder cannot remove the SplitToSequence op due to this line.

Here's a small repro of the relevant Flux module:

import torch
from diffusers import FluxTransformer2DModel

tensor_dtype = torch.bfloat16
device = "cuda"
model = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    use_safetensors=True,
    torch_dtype=tensor_dtype,
).to(device)

model=model.transformer_blocks[0].attn
s2=2
s87=4
sample_input=(
        torch.randn(s2, s87,3072, dtype=torch.bfloat16, device=device),
        torch.randn(s2, 512, 3072, dtype=torch.bfloat16, device=device),
)
dynamic_axes = (
    {0: "B", 1: "latent_dim"},
    {0: "B"},
)
torch.onnx.export(
    model,
    sample_input,
    "model.onnx",
    export_params=True,
    opset_version=19,
    dynamo=True,
    dynamic_shapes=dynamic_axes,
)

If I disable dynamic_shape then const folding works

[justinchuby]

@yuanyao-nv could you help validate the change?

[yuanyao-nv]

@justinchuby @titaiwangms Thanks for the quick fix. Yes it resolves the issue.
When is the next release of onnxscript scheduled?

cc: @asfiyab-nvidia

[justinchuby]

Next week? After our validation run over the weekend.