code: def get_transform_matrix(
      transform, trans_range, n_translate, n_rotate, n_dims
  ):
      """Generate the transformation matrix given translation and rotation.
  
      Args:
        transform: [batch_size, n_sample_frames, n_parts, transform_dims].
        [..., :n_translate] represents translations and
        [..., -n_rotat:] represents rotations.
        trans_range: float, the range of translation bound will be [-trans_range,
          trans_range].
        n_translate: int, the dimension of translation representations.
        n_rotate: int, the dimension of rotation representations.
        n_dims: int ,the dimension of points.
  
      Returns:
        matrix: tf.Tensor, [batch_size, n_frames, n_parts, (n_dims)^2].
      """
      batch_size = transform.shape[0]
      n_frames = transform.shape[1]
      n_parts = transform.shape[2]
      translation, rotation = tf.split(
          transform, [n_translate, n_rotate], axis=-1
      )
      translation = (
          tf.tanh(translation) * trans_range
      )  # Get bounded translation
  
      if n_rotate == 6:
          xy = tf.reshape(
              rotation, [batch_size, n_frames, n_parts, 3, 2]
          )
          x, y = xy[..., :]
          x = tf.nn.l2_normalize(x, axis=-1)
          z = tf.linalg.cross(x, y)
          z = tf.nn.l2_normalize(z, axis=-1)
          y = tf.linalg.cross(z, x)
          rot_mat = tf.stack([x, y, z], axis=-1)
      elif n_rotate == 2:
          rotation = tf.nn.l2_normalize(rotation, axis=-1)
          rot_mat = tf.stack(
              [
                  rotation[..., 0],
                  -rotation[..., 1],
                  rotation[..., 1],
                  rotation[..., 0],
              ],
              axis=-1,
          )
      else:
          raise ValueError("n_rotate should be either 3 or 2.")
  
      rot_mat = tf.reshape(
          rot_mat, [batch_size, n_frames, n_parts, n_dims, n_dims]
      )
      translation = tf.reshape(
          translation, [batch_size, n_frames, n_parts, n_dims, 1]
      )
      ones = tf.ones_like(translation[..., :1, :])
      zeros = tf.reshape(
          tf.zeros_like(translation),
          [batch_size, n_frames, n_parts, 1, n_dims],
      )
      up = tf.concat([rot_mat, translation], axis=-1)
      bottom = tf.concat([zeros, ones], axis=-1)
      trans_mat = tf.concat([up, bottom], axis=-2)
      return tf.reshape(
          trans_mat, [batch_size, n_frames, n_parts, (n_dims + 1) ** 2]
      )
  
file_path: graphics/tensorflow_graphics/projects/nasa/lib/model_utils.py
imports: import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: get_transform_matrix(transform, trans_range, n_translate, n_rotate, n_dims) and returns: tf