1. Create a call_torch API to make Pytorch a frontend of JAX. The call_torch API achieves two tasks:
(1) Convert the torch nn.Module or torch scriptmodel to onnx format by torch.onnx.export
(2) Use jaxonnxruntime to convert the ONNX model into JAX function
We hook all necessary components to make life easier for end PyTorch users. The input format should follow the torch.onnx.export API.
2. Improve the tests to cover more corner cases and be more robust
3. Fix more ops to support more models and tests from ONNX backend test suite