code: def test_tiled_backward(in_splits, out_splits, bias, in_f, out_f):
      base = torch.nn.Linear(in_f, out_f, bias=bias)
      test = TiledLinear(
          in_f,
          out_f,
          bias=bias,
          init_linear=copy.deepcopy(base),
          out_splits=out_splits,
          in_splits=in_splits,
      )
  
      inp = torch.rand(in_f)
  
      base_out = base(copy.deepcopy(inp))
      test_out = test(copy.deepcopy(inp))
      assert torch.allclose(base_out, test_out, rtol=1e-4)
  
      base_out.sum().backward()
      test_out.sum().backward()
  
      # compare grads
      for row in range(out_splits):
          rstart = test.out_parts[row]
          rstop = test.out_parts[row + 1]
  
          for col in range(in_splits):
              cstart = test.in_parts[col]
              cstop = test.in_parts[col + 1]
  
              local = test.linears[row][col]
              base_grad = base.weight.grad[rstart:rstop, cstart:cstop]
              assert torch.allclose(
                  base_grad, local.weight.grad, rtol=1e-4
              )
  
              if local.bias is not None:
                  base_grad = base.bias.grad[rstart:rstop]
                  assert torch.allclose(
                      base_grad, local.bias.grad, rtol=1e-4
                  )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/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 4 for calling a global function with example usage: test_tiled_backward(in_splits, out_splits, bias, in_f, out_f) without return types