Is there any way to directly execute the cached computational graph about xla HOT 5 OPEN

loss = fwd(input)
loss.backward()
optimizer.step()
graph_hash = torch_xla._XLAC._get_graph_hash([loss] + [all_parameter_gradient]) 

import torch
import torch_xla
import torch_xla.core.xla_model as xm
from torch import nn
from torch.utils._pytree import tree_map_only
from torch_xla.core.dynamo_bridge import GraphInputMatcher
from torch_xla.amp import syncfree

class M(nn.Module):

  def __init__(self):
    super().__init__()
    self.linear = nn.Linear(5, 3)

  def forward(self, x):
    return self.linear(x)

  def get_example_inputs(self):
    return (torch.rand(10, 5),)

xla_dev = xm.xla_device()
model = M().to(device=xla_dev)
optimizer = syncfree.AdamW(model.parameters(), lr=0.01)
inputs = tree_map_only(torch.Tensor, lambda x: x.to(device=xla_dev),
                        model.get_example_inputs())

xm.mark_step()
args_tensor_ids = [
    torch_xla._XLAC._xla_get_tensor_id(xla_arg) for xla_arg in inputs
]
tensor_id_to_arg_idx = {
    tensor_id: i for i, tensor_id in enumerate(args_tensor_ids)
}
output = model(*inputs).sum()
output.backward()
found_inf = torch.isnan(output).to(torch.float32).to(xla_dev)
optimizer.step(found_inf=found_inf)
opt_state = []
for name, p in model.named_parameters():
    if p.grad is not None:
        opt_state.append(p)
        opt_state.append(p.grad)
    else:
        print(name, "no grad")
output_list = [output] + opt_state
xla_graph_hash = torch_xla._XLAC._get_graph_hash(output_list)
torch_xla._XLAC._xla_warm_up_cache(output_list, [])
(
    graph_input_tensor_ids,
    graph_input_xla_values,
) = torch_xla._XLAC._get_tensors_xla_device_data_node(output_list)
xla_args_tensor_ids = set(
    tree_map_only(torch.Tensor,
                    lambda input: torch_xla._XLAC._xla_get_tensor_id(input),
                    inputs))
graph_input_matcher = GraphInputMatcher(tensor_id_to_arg_idx,
                                        graph_input_tensor_ids,
                                        graph_input_xla_values,
                                        xla_args_tensor_ids)
for i in range(3):
    graph_input = graph_input_matcher(inputs)
    res = torch_xla._XLAC._run_cached_graph(xla_graph_hash, graph_input)
    print(res[0])

import torch
import torch_xla
import torch_xla.core.xla_model as xm
from torch import nn
from torch.utils._pytree import tree_map_only
from torch_xla.core.dynamo_bridge import GraphInputMatcher
from torch_xla.amp import syncfree

class M(nn.Module):

  def __init__(self):
    super().__init__()
    self.linear = nn.Linear(5, 3)

  def forward(self, x):
    return self.linear(x)

  def get_example_inputs(self):
    return (torch.rand(10, 5),)

xla_dev = xm.xla_device()
model = M().to(device=xla_dev)
optimizer = syncfree.AdamW(model.parameters(), lr=0.01)
inputs = tree_map_only(torch.Tensor, lambda x: x.to(device=xla_dev),
                        model.get_example_inputs())

xm.mark_step()
args_tensor_ids = [
    torch_xla._XLAC._xla_get_tensor_id(xla_arg) for xla_arg in inputs
]
tensor_id_to_arg_idx = {
    tensor_id: i for i, tensor_id in enumerate(args_tensor_ids)
}
for i in range(3):
    output = model(*inputs).sum()
    output.backward()
    found_inf = torch.isnan(output).to(torch.float32).to(xla_dev)
    optimizer.step(found_inf=found_inf)
    optimizer.zero_grad()
    xm.mark_step()
    print("debug ", output)