code: def get_locks(size, dev):
      if dev not in _sparse_matmul.locks or size > _sparse_matmul.locks[
          dev
      ].size(0):
          _sparse_matmul.locks[dev] = torch.zeros(
              size, dtype=torch.int32, device=dev
          )
      return _sparse_matmul.locks[dev]
  
imports: import importlib
  import torch
  import triton
  import triton.language
  import triton._C.libtriton
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 2 for calling the _sparse_matmul obj's underlying member method code with example usage: obj.get_locks(size, dev) and returns: _sparse_matmul