Writing a compiler for PyTorch Models

I am currently working with a custom accelerator for machine learning models. Currently models are trained in PyTorch to get the weights and then implemented using a c driver for the accelerator. I am interested in writing a compiler that consumes a PyTorch model and outputs driver calls. My question is whether torch script or another PyTorch IR provides a sufficient API that exposes the tensor operations so I can iterate over them into function calls. Any advice about feasibility or implementation would be highly useful.