code: def test_repeat_interleave():
      patch_fn = patched_function.torch_repeat_interleave
  
      # examples from https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html
      data = torch.tensor([1, 2, 3])
      materialized_output = torch.repeat_interleave(data, repeats=2)
      repeat_interleave = partial(patch_fn, repeats=2)
      meta_data = data.to("meta")
      _assert_output_shape(
          data=meta_data,
          patch_fn=repeat_interleave,
          expect_exception=False,
          output_shape=materialized_output.shape,
      )
  
      data = torch.tensor([[1, 2], [3, 4]])
      materialized_output = torch.repeat_interleave(
          data, repeats=3, dim=1
      )
      repeat_interleave = partial(patch_fn, repeats=3, dim=1)
      meta_data = data.to("meta")
      _assert_output_shape(
          data=meta_data,
          patch_fn=repeat_interleave,
          expect_exception=False,
          output_shape=materialized_output.shape,
      )
  
      data = torch.tensor([[1, 2], [3, 4]])
      materialized_output = torch.repeat_interleave(
          data, repeats=torch.tensor([1, 2]), dim=-1
      )
      repeat_interleave = partial(
          patch_fn, repeats=torch.tensor([1, 2]), dim=-1
      )
      meta_data = data.to("meta")
      _assert_output_shape(
          data=meta_data,
          patch_fn=repeat_interleave,
          expect_exception=False,
          output_shape=materialized_output.shape,
      )
  
      data = torch.tensor([[1, 2], [3, 4]])
      materialized_output = torch.repeat_interleave(
          data, repeats=torch.tensor([1, 2]), dim=0
      )
      repeat_interleave = partial(patch_fn, repeats=[1, 2], dim=0)
      meta_data = data.to("meta")
      _assert_output_shape(
          data=meta_data,
          patch_fn=repeat_interleave,
          expect_exception=True,
          output_shape=materialized_output.shape,
      )
  
file_path: ColossalAI/tests/test_fx/test_tracer/test_patched_op.py
imports: from functools import partial
  import torch
  from colossalai.fx.tracer.meta_patch import patched_function
  from colossalai.testing import clear_cache_before_run
label_desc: Use image node_id 3 for calling a global function with example usage: test_repeat_interleave() without return types