code: def test_tiled_init(in_splits, out_splits):
      in_f = 32
      out_f = 40
      base = torch.nn.Linear(in_f, out_f, bias=True)
      l = TiledLinear(
          in_f,
          out_f,
          bias=True,
          init_linear=copy.deepcopy(base),
          out_splits=out_splits,
          in_splits=in_splits,
      )
  
      for out_id in range(out_splits):
          for in_id in range(in_splits):
              local_l = l.linears[out_id][in_id]
              assert isinstance(local_l, torch.nn.Linear)
  
              rstart = l.out_parts[out_id]
              rstop = l.out_parts[out_id + 1]
              cstart = l.in_parts[in_id]
              cstop = l.in_parts[in_id + 1]
  
              local_out = rstop - rstart
              local_in = cstop - cstart
              assert (
                  local_l.weight.size()[1] == local_in
              ), f"local[{out_id}][{in_id}].size {local_l.weight.size()}"
              assert local_l.weight.size()[0] == local_out
  
              test = base.weight[rstart:rstop, cstart:cstop]
  
              assert local_l.weight.size() == test.size()
              assert torch.equal(local_l.weight.data, test.data)
  
              if in_id == in_splits - 1:
                  assert local_l.bias is not None
                  assert local_l.bias.size()[0] == local_out
              else:
                  assert local_l.bias is None
  
file_path: DeepSpeed/tests/unit/runtime/zero/test_zero_tiled.py
imports: import copy
  import torch
  from deepspeed.runtime.zero.tiling import TiledLinear, TiledLinearReturnBias
  import pytest
label_desc: Use image node_id 1 for calling a global function with example usage: test_tiled_init(in_splits, out_splits) without return types