code: def im2row_index(
      input_shape,
      block_shape,
      rank=2,
      slice_step=(1, 1),
      dilations=(1, 1),
      dtype=tf.int32,
      transpose=False,
      validate_args=False,
      name=None,
  ):
      """Computes indexes into a flattened image for building `im2row`."""
      with tf.name_scope(name or "im2row_index"):
          if tf.get_static_value(rank) != 2:
              raise NotImplementedError(
                  "Argument `rank` currently only supports `2`; "
                  'saw "{}".'.format(rank)
              )
          fh, fw = prepare_tuple_argument(
              block_shape,
              n=rank,
              arg_name="block_shape",
              validate_args=validate_args,
          )
          sh, sw = prepare_tuple_argument(
              slice_step,
              n=rank,
              arg_name="slice_step",
              validate_args=validate_args,
          )
          dh, dw = prepare_tuple_argument(
              dilations,
              n=rank,
              arg_name="dilations",
              validate_args=validate_args,
          )
  
          # 1) Process input arguments.
          batch_shape, h, w, c = ps.split(
              ps.reshape(ps.cast(input_shape, dtype=dtype), shape=[-1]),
              num_or_size_splits=[-1, 1, 1, 1],
          )
          h, w, c = h[0], w[0], c[0]
  
          tot_fh = dh * (fh - 1) + 1
          tot_fw = dw * (fw - 1) + 1
  
          # 2) Assemble all block start positions as indexes into the flattened image.
          # start_idx.shape = [fh, fw, c]
          if transpose:
              last_element = (
                  lambda size, step: size - (size - 1) % step - 1
              )
              w_step = c * dw
              h_step = c * w * dh
              last_w = last_element(c * tot_fw, w_step)
              last_h = last_element(c * w * tot_fh, h_step)
              start_idx = cartesian_add(
                  [
                      ps.range(last_h, -1, delta=-h_step, dtype=dtype),
                      ps.range(last_w, -1, delta=-w_step, dtype=dtype),
                      ps.range(c, delta=1, dtype=dtype),
                  ]
              )
          else:
              start_idx = cartesian_add(
                  [
                      ps.range(
                          c * w * tot_fh, delta=c * w * dh, dtype=dtype
                      ),
                      ps.range(c * tot_fw, delta=c * dw, dtype=dtype),
                      ps.range(c, delta=1, dtype=dtype),
                  ]
              )
  
          # 3) Assemble all block offsets (into flattened image).
          eh = h - tot_fh + 1
          ew = w - tot_fw + 1
  
          offset_idx = cartesian_add(
              [
                  ps.range(w * eh, delta=w * sh, dtype=dtype),
                  ps.range(ew, delta=sw, dtype=dtype),
              ]
          )
  
          offset_idx = offset_idx * c
          oh = (eh - 1) // sh + 1  # out height
          ow = (ew - 1) // sw + 1  # out width
  
          # 4) Combine block start/offset pairs.
          # shape = [(eh // sh) * (ew // sw), fh * fw * c]
          idx = cartesian_add([offset_idx, start_idx])
          new_shape = ps.concat(
              [
                  batch_shape,
                  ps.convert_to_shape_tensor([oh, ow, fh * fw * c]),
              ],
              axis=0,
          )
          return idx, new_shape
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.nn.util import utils
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 2 for calling a global function with example usage: im2row_index(input_shape, block_shape, rank, slice_step, dilations, dtype, transpose, validate_args, name) and returns: idx,  new_shape