Neuron import hook (#5429)

* Enable Neuron import hook for calling initialization functions if using AWS Neuron

* removing copy/paste error

* moving aws init call and removing comment

Author: aws-kingrj

torch_xla/__init__.py

@@ -60,6 +60,15 @@ def _summarize_fn_tracker():
   os.remove(_tmp_fname)
 
 
+def _aws_ec2_inf_trn_init():
+  try:
+    from torch_neuronx import xla
+  except ImportError:
+    return
+  else:
+    xla.init()
+
+
 def _setup_tpu_vm_library_path() -> bool:
   """Returns true if $TPU_LIBRARY is set or can be inferred.
 
@@ -105,6 +114,9 @@ def _setup_tpu_vm_library_path() -> bool:
 
 _found_libtpu = _setup_tpu_vm_library_path()
 
+# Setup Neuron library for AWS EC2 inf/trn instances.
+_aws_ec2_inf_trn_init()
+
 
 def _prepare_to_exit():
   _XLAC._prepare_to_exit()