code: def rot_from_axisangle(vec):
      """Convert an axisangle rotation into a 4x4 transformation matrix
      (adapted from https://github.com/Wallacoloo/printipi)
      Input 'vec' has to be Bx1x3
      """
      angle = mx.nd.norm(vec, 2, 2, keepdims=True)
      axis = vec / (angle + 1e-7)
  
      ca = mx.nd.cos(angle)
      sa = mx.nd.sin(angle)
      C = 1 - ca
  
      x = axis[..., 0].expand_dims(axis=1)
      y = axis[..., 1].expand_dims(axis=1)
      z = axis[..., 2].expand_dims(axis=1)
  
      xs = x * sa
      ys = y * sa
      zs = z * sa
      xC = x * C
      yC = y * C
      zC = z * C
      xyC = x * yC
      yzC = y * zC
      zxC = z * xC
  
      elem1 = mx.nd.squeeze(x * xC + ca).reshape((vec.shape[0], 1))
      elem2 = mx.nd.squeeze(xyC - zs).reshape((vec.shape[0], 1))
      elem3 = mx.nd.squeeze(zxC + ys).reshape((vec.shape[0], 1))
      elem4 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem5 = mx.nd.squeeze(xyC + zs).reshape((vec.shape[0], 1))
      elem6 = mx.nd.squeeze(y * yC + ca).reshape((vec.shape[0], 1))
      elem7 = mx.nd.squeeze(yzC - xs).reshape((vec.shape[0], 1))
      elem8 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem9 = mx.nd.squeeze(zxC - ys).reshape((vec.shape[0], 1))
      elem10 = mx.nd.squeeze(yzC + xs).reshape((vec.shape[0], 1))
      elem11 = mx.nd.squeeze(z * zC + ca).reshape((vec.shape[0], 1))
      elem12 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem13 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem14 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem15 = (
          mx.nd.zeros(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
      elem16 = (
          mx.nd.ones(vec.shape[0])
          .as_in_context(context=vec.context)
          .reshape((vec.shape[0], 1))
      )
  
      elem = [
          elem1,
          elem2,
          elem3,
          elem4,
          elem5,
          elem6,
          elem7,
          elem8,
          elem9,
          elem10,
          elem11,
          elem12,
          elem13,
          elem14,
          elem15,
          elem16,
      ]
      rot = mx.nd.concat(*elem, dim=1).reshape((vec.shape[0], 4, 4))
  
      return rot
  
imports: from __future__ import absolute_import, division, print_function
  import numpy
  import mxnet
  from mxnet.gluon import nn
label_desc: Use image node_id 4 for calling a global function with example usage: rot_from_axisangle(vec) and returns: rot