code: def block_sizes_strategy(
      draw: Draw,
      q_seq_len: int,
      kv_seq_len: int,
      include_bwd_blocks: bool = False,
      use_fused_bwd_kernel: bool = False,
  ) -> splash.BlockSizes:
      all_block_shapes = [128, 256, 512]
      q_layout = draw(hps.sampled_from(splash.QKVLayout))
      k_layout = draw(hps.sampled_from(splash.QKVLayout))
      v_layout = draw(hps.sampled_from(splash.QKVLayout))
      layouts = dict(
          q_layout=q_layout, k_layout=k_layout, v_layout=v_layout
      )
      q_valid_block_shapes = [
          bs for bs in all_block_shapes if bs <= q_seq_len
      ]
      kv_valid_block_shapes = [
          bs for bs in all_block_shapes if bs <= kv_seq_len
      ]
      bq, bkv = (
          draw(hps.sampled_from(q_valid_block_shapes)),
          draw(hps.sampled_from(kv_valid_block_shapes)),
      )
      bkv_compute = draw(
          hps.sampled_from(
              [None, *[b for b in kv_valid_block_shapes if b <= bkv]]
          )
      )
      if not include_bwd_blocks:
          return splash.BlockSizes(
              block_q=bq,
              block_kv=bkv,
              block_kv_compute=bkv_compute,
              **layouts,
          )
      all_block_shapes = [128, 256]
      q_valid_block_shapes = [
          bs for bs in all_block_shapes if bs <= q_seq_len
      ]
      kv_valid_block_shapes = [
          bs for bs in all_block_shapes if bs <= kv_seq_len
      ]
      bq_dkv, bkv_dkv = (
          draw(hps.sampled_from(q_valid_block_shapes)),
          draw(hps.sampled_from(kv_valid_block_shapes)),
      )
      if use_fused_bwd_kernel:
          bq_dq, bkv_dq = None, None
      else:
          bq_dq = draw(hps.sampled_from(q_valid_block_shapes))
          bkv_dq = draw(hps.sampled_from(kv_valid_block_shapes))
      block_kv_dkv_compute = draw(
          hps.sampled_from(
              [
                  None,
                  *[b for b in kv_valid_block_shapes if b <= bkv_dkv],
              ]
          )
      )
      return splash.BlockSizes(
          block_q=bq,
          block_kv=bkv,
          block_kv_compute=bkv_compute,
          block_q_dkv=bq_dkv,
          block_kv_dkv=bkv_dkv,
          block_kv_dkv_compute=block_kv_dkv_compute,
          block_q_dq=bq_dq,
          block_kv_dq=bkv_dq,
          use_fused_bwd_kernel=use_fused_bwd_kernel,
          **layouts,
      )
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, Callable, TypeVar
  from absl.testing import absltest
  from absl.testing import parameterized
  import hypothesis
  import hypothesis.strategies
  import jax
  from jax import random
  from jax import tree_util
  from jax._src import test_util
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  import jax.numpy
  import numpy
label_desc: Use image node_id 13 for calling a global function with example usage: block_sizes_strategy(draw, q_seq_len, kv_seq_len, include_bwd_blocks, use_fused_bwd_kernel) and returns: splash, splash