code: def ref_affine_grid(theta, size, align_corners):
      if len(size) == 2:
          grid_s = affine_grid_2d(theta, size, align_corners)
      elif len(size) == 3:
          grid_s = affine_grid_3d(theta, size, align_corners)
      return grid_s
  
file_path: nnabla/python/test/function/test_affine_grid.py
imports: import pytest
  import numpy
  import nnabla
  import nnabla.functions
  from nbla_test_utils import list_context
  from refs import generate_transformation_2d, generate_transformation_3d, affine_grid_2d, affine_grid_3d
label_desc: Use image node_id 2 for calling a global function with example usage: ref_affine_grid(theta, size, align_corners) and returns: grid_s