code: def transform_tf(
      x_e: Tensor,
      translation: Tuple[float, float] = (0.0, 0.0),
      rotation: float = 0.0,
  ) -> Any:
      """
      Input
      - x: Ep tensor of shape (bs, n_x, n_y, C).
      - translation: tuple of x, y translation in pixels
      - rotation: rotation in rad
  
      Returns
      - out_fmap: transformed input feature map. Tensor of size (bs, n_x, n_y, C).
      Notes
  
      References:
      [#Jade]: 'Spatial Transformer Networks', Jaderberg et. al,
           (https://arxiv.org/abs/1506.02025)
      """
      import tensorflow as tf
  
      bs = x_e.shape[0]
      theta = np.zeros((2, 3)).astype(np.float32)
  
      theta[0, :] = [
          np.cos(rotation),
          -np.sin(rotation),
          translation[0],
      ]
      theta[1, :] = [np.sin(rotation), np.cos(rotation), translation[1]]
      theta = np.tile(theta[None], (bs, 1, 1)).reshape(bs, 2, 3)
      # convert from pixels to relative translation (bs, x, y, n_ch)
      theta[:, 0, 2] /= x_e.shape[1] / 2.0
      theta[:, 1, 2] /= x_e.shape[2] / 2.0
  
      # to tf
      theta = tf.convert_to_tensor(theta)
      x = x_e.raw
  
      # grab input dimensions
      assert theta.shape[1:] == (2, 3)
      assert len(x.shape) == 4
      bs = tf.shape(x)[0]
      n_x = tf.shape(x)[1]  # height matrix
      n_y = tf.shape(x)[2]  # width matrix
  
      def get_pixel_value(img: Any, x: Any, y: Any) -> Any:
          """
          Utility function to get pixel value for coordinate
          vectors x and y from a  4D tensor image.
  
          Args:
          - img: tensor of shape (bs, n_x, n_y, C)
          - x: flattened tensor of shape (bs*n_x*n_y,)
          - y: flattened tensor of shape (bs*n_x*n_y,)
  
          Returns:
          - output: tensor of shape (bs, n_x, n_y, C)
          """
          batch_idx = tf.range(0, bs)
          batch_idx = tf.reshape(batch_idx, (bs, 1, 1))
          b = tf.tile(batch_idx, (1, n_x, n_y))
          indices = tf.stack([b, y, x], 3)
          return tf.gather_nd(img, indices)
  
      def bilinear_sampler(img: Any, x: Any, y: Any) -> Any:
          """
          Performs bilinear sampling of the input images according to the
          normalized coordinates provided by the sampling grid. Note that
          the sampling is done identically for each channel of the input.
          To test if the function works properly, output image should be
          identical to input image when theta is initialized to identity
          transform.
  
          Args:
          - img: batch of images in (bs, n_x, n_y, C) layout.
          - grid: x, y which is the output of affine_grid_generator.
  
          Returns:
          - out: interpolated images according to grids. Same size as grid.
          """
          max_y = tf.cast(n_x - 1, "int32")
          max_x = tf.cast(n_y - 1, "int32")
  
          # rescale x and y to [0, n_y-1/n_x-1]
          x = tf.cast(x, "float32")
          y = tf.cast(y, "float32")
          x = 0.5 * ((x + 1.0) * tf.cast(max_x, "float32"))
          y = 0.5 * ((y + 1.0) * tf.cast(max_y, "float32"))
  
          # grab 4 nearest corner points for each (x_i, y_i)
          x0 = tf.cast(tf.floor(x), "int32")
          x1 = x0 + 1
          y0 = tf.cast(tf.floor(y), "int32")
          y1 = y0 + 1
  
          # clip to range [0, n_x-1/n_y-1] to not violate img boundaries
          min_val = 0
          x0 = tf.clip_by_value(x0, min_val, max_x)
          x1 = tf.clip_by_value(x1, min_val, max_x)
          y0 = tf.clip_by_value(y0, min_val, max_y)
          y1 = tf.clip_by_value(y1, min_val, max_y)
  
          # get pixel value at corner coords
          Ia = get_pixel_value(img, x0, y0)
          Ib = get_pixel_value(img, x0, y1)
          Ic = get_pixel_value(img, x1, y0)
          Id = get_pixel_value(img, x1, y1)
  
          # recast as float for delta calculation
          x0 = tf.cast(x0, "float32")
          x1 = tf.cast(x1, "float32")
          y0 = tf.cast(y0, "float32")
          y1 = tf.cast(y1, "float32")
  
          # calculate deltas
          wa = (x1 - x) * (y1 - y)
          wb = (x1 - x) * (y - y0)
          wc = (x - x0) * (y1 - y)
          wd = (x - x0) * (y - y0)
  
          # add dimension for addition
          wa = tf.expand_dims(wa, axis=3)
          wb = tf.expand_dims(wb, axis=3)
          wc = tf.expand_dims(wc, axis=3)
          wd = tf.expand_dims(wd, axis=3)
  
          # compute output
          out = tf.add_n([wa * Ia, wb * Ib, wc * Ic, wd * Id])
  
          return out
  
      def affine_grid_generator(
          height: Any, width: Any, theta: Any
      ) -> Any:
          """
          This function returns a sampling grid, which when
          used with the bilinear sampler on the input feature
          map, will create an output feature map that is an
          affine transformation [1] of the input feature map.
  
          Args:
          - height: desired height of grid/output. Used
            to downsample or upsample.
          - width: desired width of grid/output. Used
            to downsample or upsample.
          - theta: affine transform matrices of shape (num_batch, 2, 3).
            For each image in the batch, we have 6 theta parameters of
            the form (2x3) that define the affine transformation T.
  
          Returns:
          - normalized grid (-1, 1) of shape (num_batch, 2, n_x, n_y).
            The 2nd dimension has 2 components: (x, y) which are the
            sampling points of the original image for each point in the
            target image.
          Note
          ----
          [1]: the affine transformation allows cropping, translation,
               and isotropic scaling.
          """
          num_batch = tf.shape(theta)[0]
  
          # create normalized 2D grid
          x_l = tf.linspace(-1.0, 1.0, width)
          y_l = tf.linspace(-1.0, 1.0, height)
          x_t, y_t = tf.meshgrid(x_l, y_l)
  
          # flatten
          x_t_flat = tf.reshape(x_t, [-1])
          y_t_flat = tf.reshape(y_t, [-1])
  
          # reshape to [x_t, y_t , 1] - (homogeneous form)
          ones = tf.ones_like(x_t_flat)
          sampling_grid = tf.stack([x_t_flat, y_t_flat, ones])
  
          # repeat grid num_batch times
          sampling_grid = tf.expand_dims(sampling_grid, axis=0)
          sampling_grid = tf.tile(
              sampling_grid, tf.stack([num_batch, 1, 1])
          )
  
          # cast to float32 (required for matmul)
          theta = tf.cast(theta, "float32")
          sampling_grid = tf.cast(sampling_grid, "float32")
  
          # transform the sampling grid - batch multiply
          batch_grids = tf.matmul(theta, sampling_grid)
          # batch grid has shape (num_batch, 2, n_x*n_y)
  
          # reshape to (num_batch, n_x, n_y, 2)
          batch_grids = tf.reshape(
              batch_grids, [num_batch, 2, height, width]
          )
  
          return batch_grids
  
      batch_grids = affine_grid_generator(n_x, n_y, theta)
  
      x_s = batch_grids[:, 0, :, :]
      y_s = batch_grids[:, 1, :, :]
  
      # sample input with grid to get output
      transformed_images = bilinear_sampler(x, x_s, y_s)
  
      return astensor(transformed_images)
  
imports: from typing import Tuple, Any
  import numpy
  import math
  from eagerpy import astensor, Tensor
  from eagerpy.tensor import TensorFlowTensor, PyTorchTensor
label_desc: Use image node_id 3 for calling a global function with example usage: transform_tf(x_e, translation, rotation) and returns: astensor, tf, out, batch_grids