code: def segment_mask(
      q_segment_ids: jax.Array,
      kv_segment_ids: jax.Array,
  ):
      # [B, T, 1] or [T, 1]
      q_segment_ids = jnp.expand_dims(q_segment_ids, axis=-1)
      # [B, 1, S] or [1, S]
      if kv_segment_ids.ndim == 1:
          kv_segment_ids = jnp.expand_dims(kv_segment_ids, axis=0)
      else:
          kv_segment_ids = jnp.expand_dims(kv_segment_ids, axis=1)
      return jnp.equal(q_segment_ids, kv_segment_ids).astype(jnp.bool_)
  
file_path: jax/jax/experimental/pallas/ops/attention.py
imports: from __future__ import annotations
  import functools
  from typing import Any, Optional
  import jax
  from jax import lax
  from jax.experimental import pallas
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: segment_mask(q_segment_ids, kv_segment_ids) and returns: jnp