code: def make_layout(rho, shape):
      probs = torch.Tensor([rho, 1 - rho])
      generator = torch.distributions.categorical.Categorical(probs)
      layout = generator.sample(shape)
      return layout
  
file_path: DeepSpeed/tests/unit/ops/sparse_attention/test_sparse_attention.py
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import SparseAttnBuilder
  from unit.util import skip_on_arch, skip_on_cuda
label_desc: Use image node_id 4 for calling a global function with example usage: make_layout(rho, shape) and returns: layout