code: def actor_prediction(
      model: AlphaTensorModel,
      input_tensor: torch.Tensor,
      maximum_rank: int,
      mc_n_sim: int,
      N_bar: int,
      return_actions: bool = False,
  ):
      """Runs the monte carlo tree search algorithm to obtain the next states,
      policies and rewards.
  
      Args:
          model (AlphaTensorModel): The model to use for the simulation.
          input_tensor (torch.Tensor): The initial state.
          maximum_rank (int): The maximum number of steps to simulate.
          mc_n_sim (int): The number of simulations to run.
          N_bar (int): The parameter used to compute the improved policy.
          return_actions (bool): If True, only actions are returned.
      """
      # input_tensor has shape (1, T, S, S, S)
      state = input_tensor
      rank = 0
      game_tree = {}
      state_dict = {}
      hash_states = []
      states = []
      while rank < maximum_rank:
          states.append(state)
          hash_states.append(to_hash(extract_present_state(state)))
          state = monte_carlo_tree_search(
              model,
              state,
              mc_n_sim,
              rank,
              maximum_rank,
              game_tree,
              state_dict,
          )
          if game_is_finished(extract_present_state(state)):
              break
          rank += 1
      final_state = extract_present_state(state)
      policies = compute_improved_policy(
          state_dict, hash_states, model.n_steps, model.n_logits, N_bar
      )
      reward = (
          int(torch.linalg.matrix_rank(final_state).sum())
          if not game_is_finished(final_state)
          else 0
      )
      rewards = torch.cumsum(
          torch.tensor([-1] * (len(policies) - 1) + [reward]), dim=0
      )
      if return_actions:
          actions = [
              state_dict[hash_state][5] for hash_state in hash_states
          ]
          return actions
      # policies do not have the batch size, but states still have it
      states = [s.squeeze(0) for s in states]
      return states, policies, rewards
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/actors/stage.py
imports: from typing import Dict, List
  import torch
  from open_alpha_tensor.core.data.utils import get_scalars, map_action_to_triplet
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 15 for calling a global function with example usage: actor_prediction(model, input_tensor, maximum_rank, mc_n_sim, N_bar, return_actions) and returns: states,  policies,  rewards, actions