code: def _next_nonzero(
      h,
      i,
      j,
      data_next_ref,
      block_mask_ref,
      m_next_ref,
      next_i=False,
  ):
      assert (data_next_ref is None) == (block_mask_ref is None)
  
      if data_next_ref is None and block_mask_ref is None:
          # Handle the case in which we have no masking nor next data information.
          # Simply fetch the next data and apply the mask for every block.
          assert m_next_ref is None
          next_data = i if next_i else j
          return (
              next_data,
              None,  # next mask
              True,  # should run
              False,  # should not mask
          )
  
      assert data_next_ref.shape == block_mask_ref.shape
      assert (
          m_next_ref is None
          or data_next_ref.shape[0] == m_next_ref.shape[0]
      )
  
      # We are working with one head only. Force the head index to 0.
      if data_next_ref.shape[0] == 1:
          h = 0
  
      # When scalar-memory data is of types smaller than int32, then we have to
      # upcast it back to use it in the kernel.
  
      to_i32 = lambda x: x.astype(jnp.int32)
  
      is_nonzero = to_i32(block_mask_ref[h, i, j]) > 0
      if m_next_ref is None:
          should_not_mask = True
          next_m = None
      else:
          should_not_mask = to_i32(block_mask_ref[h, i, j]) != 1
          next_m = to_i32(m_next_ref[h, i, j])
      next_j = to_i32(data_next_ref[h, i, j])
      return next_j, next_m, is_nonzero, should_not_mask
  
imports: from __future__ import annotations
  import dataclasses
  import enum
  import functools
  from typing import Any, Callable, Literal, NamedTuple, Union, Optional, overload
  import jax
  from jax import ad_checkpoint
  from jax import lax
  from jax import tree_util
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info
  import jax.numpy
  import numpy
label_desc: Use image node_id 12 for calling a global function with example usage: _next_nonzero(h, i, j, data_next_ref, block_mask_ref, m_next_ref, next_i) and returns: next_j,  next_m,  is_nonzero,  should_not_mask, next_data,  None,  True,  False