code: def below_or_on_diag(r, r_blk_size, c, c_blk_size):
      # A block is considered below or on diagonal as long as the bottom left
      # corner of the block is below or on diagonal.
      return ((r + 1) * r_blk_size - 1) > (c * c_blk_size)
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, NamedTuple
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 5 for calling a global function with example usage: below_or_on_diag(r, r_blk_size, c, c_blk_size) and returns: unknown