Thanks, if removing bailout is possible that would definitely work for me. At a quick glance there is no pass to remove Bailout nodes in torch, but I’ll try if I can do it from “userland”.
Another API I’m interested in is _propagate_shapes function:

  
      github.com
  
  
    pytorch/pytorch/blob/8527ba8b7059d8c7df6fc2c11271ceefbf91ce5a/torch/csrc/jit/script/init.cpp#L476-L485

static std::shared_ptr<Graph> _propagate_shapes(
    Graph& graph,
    std::vector<at::Tensor> inputs,
    bool with_grad = false) {
  Stack stack(inputs.begin(), inputs.end());
  auto retval = graph.copy();
  setInputTensorTypes(*retval, stack, /*complete=*/false);
  PropagateInputShapes(retval);
  return retval;
}



  
  
    
    
  
  


Do you know the use case of this function and if this is something I should take a look?