code: def compute_start_indices(self, loop_idx, *args):
      (
          discharged_jaxpr,
          discharged_consts,
      ) = state_discharge.discharge_state(
          self.index_map_jaxpr.jaxpr, self.index_map_jaxpr.consts
      )
      jaxpr = jax_core.ClosedJaxpr(discharged_jaxpr, discharged_consts)
      block_indices_and_rest = jax_core.jaxpr_as_fun(jaxpr)(
          *loop_idx, *args
      )
      # Since we're passing in `Ref`s potentially, we need to split out their
      # updated values since we only care about the return values.
      block_indices, _ = split_list(
          block_indices_and_rest, [len(self.block_shape)]
      )
      return tuple(
          i if b is mapped else b * i
          for b, i in zip(self.block_shape, block_indices)
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import functools
  from typing import Any, Callable
  from collections.abc import Iterator
  from jax._src import api_util
  from jax._src import core
  from jax._src import linear_util
  from jax._src import state
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import partial_eval
  from jax._src.state import discharge
  import jax.numpy
label_desc: Use image node_id 1 for calling the BlockMapping obj's underlying member method code with example usage: obj.compute_start_indices(loop_idx) and returns: tuple