code: def create_inputs(rng, batch_size, size):
      if len(size) == 2:
          theta = generate_transformation_2d(rng, batch_size)
      elif len(size) == 3:
          theta = generate_transformation_3d(rng, batch_size)
      return theta
  
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 1 for calling a global function with example usage: create_inputs(rng, batch_size, size) and returns: theta