Improves implementation of aten_index_put

[xadupre]

From #2606.

[codecov]

Codecov Report

❌ Patch coverage is 75.26882% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.46%. Comparing base (8a94ad6) to head (1c1fc1a).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	75.26%	13 Missing and 10 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2641   +/-   ##
=======================================
  Coverage   70.46%   70.46%           
=======================================
  Files         224      224           
  Lines       26572    26657   +85     
  Branches     2637     2654   +17     
=======================================
+ Hits        18723    18785   +62     
- Misses       6928     6941   +13     
- Partials      921      931   +10     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.