code: def get_step_idx(state: TreeState) -> tf.Tensor:
      """Returns the current leaf node index based on `TreeState.level_buffer_idx`."""
      step_idx = tf.constant(-1, dtype=tf.int32)
      for i in tf.range(len(state.level_buffer_idx)):
          step_idx += tf.math.pow(2, state.level_buffer_idx[i])
      return step_idx
  
imports: import abc
  import collections
  from typing import Any, Callable, Collection, NamedTuple, Optional, Tuple, Union
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: get_step_idx(state) and returns: step_idx