jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def vectors2rotmat(vec_rest, vec_final, dtype):
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc
//  [python]
def vectors2rotmat(vec_rest, vec_final, dtype):
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.linalg.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.linalg.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc
//  [python]
def mpjpe_cal_mask(predicted, target, mask):
    assert predicted.shape == target.shape
    # index = [i for i in range(17) if i in mask]
    predicted = predicted[:,:,mask,:]
    target = target[:,:,mask,:]
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1)).contiguous()
//  [python]
def mpjpe_cal_mask(predicted, target, mask):
    assert predicted.shape == target.shape
    # index = [i for i in range(17) if i in mask]
    predicted = predicted[:,:,mask,:]
    target = target[:,:,mask,:]
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1)).contiguous()
//  [python]
def forward(self, states):
        # Generate expected 2d landmark locations from the current state and 3d landmark locations
        # Input: states: N*state_dim
        # Output: expected 2d landmark x coordinate: N*N_landmarks
        #         expected 2d landmark y coordinate: N*N_landmarks
        #        expected 2d landmark x covariance: N*N_landmarks*N_landmarks
        #        expected 2d landmark y covariance: N*N_landmarks*N_landmarks
        N, state_dim = states.shape
        self.state_dim = state_dim
        
        # Extract the orientation quaternion from state variables
        quat = states[:, 6:10].detach().clone()
        
        # Normalize the quaternion (world -> body)
        orientation = torch.div(quat, torch.norm(quat, dim=1)[:, None])
        # Retrieve the landmark positions in previous frame
        landmarks = self.landmarks  # N_landmarks*3
        
        # Previous quaternion
        prev_quat = self.prev_state[:, 6:10].detach().clone()
        
        # Compute the change in rotation between frames (TODO: unit test this)
        # delta_quat = quat_delta(prev_quat, quat)  # N*3
        delta_quat = quat_delta(quat, quat)
        
        # Compute the ENU motion from state variables
        vel_enu = (states[:, :3]-self.prev_state[:, :3])
        # Calculate expected 2d landmark locations
        img_pts = self._vel_to_2dfeatures(vel_enu, orientation, delta_quat)
        
        # Calculate covariance of expected 2d landmark locations
        R = self.robust_cost_cholesky(self.linearization_point, img_pts, vel_enu, orientation, delta_quat)   # N*(N_landmarks*2)*(N_landmarks*2)

        return img_pts.float(), R.float()
//  [python]
def forward(self, states):
        # Generate expected 2d landmark locations from the current state and 3d landmark locations
        # Input: states: N*state_dim
        # Output: expected 2d landmark x coordinate: N*N_landmarks
        #         expected 2d landmark y coordinate: N*N_landmarks
        #        expected 2d landmark x covariance: N*N_landmarks*N_landmarks
        #        expected 2d landmark y covariance: N*N_landmarks*N_landmarks
        N, state_dim = states.shape
        self.state_dim = state_dim
        
        # Extract the orientation quaternion from state variables
        quat = states[:, 6:10].detach().clone()
        
        # Normalize the quaternion (world -> body)
        orientation = torch.div(quat, torch.linalg.norm(quat, dim=1)[:, None])
        # Retrieve the landmark positions in previous frame
        landmarks = self.landmarks  # N_landmarks*3
        
        # Previous quaternion
        prev_quat = self.prev_state[:, 6:10].detach().clone()
        
        # Compute the change in rotation between frames (TODO: unit test this)
        # delta_quat = quat_delta(prev_quat, quat)  # N*3
        delta_quat = quat_delta(quat, quat)
        
        # Compute the ENU motion from state variables
        vel_enu = (states[:, :3]-self.prev_state[:, :3])
        # Calculate expected 2d landmark locations
        img_pts = self._vel_to_2dfeatures(vel_enu, orientation, delta_quat)
        
        # Calculate covariance of expected 2d landmark locations
        R = self.robust_cost_cholesky(self.linearization_point, img_pts, vel_enu, orientation, delta_quat)   # N*(N_landmarks*2)*(N_landmarks*2)

        return img_pts.float(), R.float()
//  [python]
def frame_loss(predicted):#256,9,17,2
    loss = 0
    for k in range(predicted.size(0)-1):
        for i in range(predicted.size(1)-1):
            for j in range(predicted.size(2)-1):
                loss += (predicted[k+1,i+1,j+1,0] - predicted[k,i,j,0])**2
                loss += (predicted[k+1,i+1,j+1,1] - predicted[k,i,j,1])**2
    return loss

    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1)).contiguous()
//  [python]
def frame_loss(predicted):#256,9,17,2
    loss = 0
    for k in range(predicted.size(0)-1):
        for i in range(predicted.size(1)-1):
            for j in range(predicted.size(2)-1):
                loss += (predicted[k+1,i+1,j+1,0] - predicted[k,i,j,0])**2
                loss += (predicted[k+1,i+1,j+1,1] - predicted[k,i,j,1])**2
    return loss

    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1)).contiguous()
//  [python]
def skeloss(predicted, target):
    assert predicted.shape == target.shape
    start = [0,1,2, 0,4,5, 0,7,8,9,   8,14,15,  8,11,12]
    end =   [1,2,3, 4,5,6, 7,8,9,10, 14,15,16, 11,12,13]
    ske_predicted = torch.zeros(len(start))
    ske_target = torch.zeros(len(start))
    for i in range(len(start)):
        ske_predicted[i] = torch.mean(torch.norm(predicted[:,:,start[i],:] - predicted[:,:,end[i],:], dim=2)).contiguous()
        ske_target[i] = torch.mean(torch.norm(target[:,:,start[i],:] - target[:,:,end[i],:], dim=2)).contiguous()



    return torch.mean(torch.norm(ske_predicted[i] - ske_target[i]))
//  [python]
def skeloss(predicted, target):
    assert predicted.shape == target.shape
    start = [0,1,2, 0,4,5, 0,7,8,9,   8,14,15,  8,11,12]
    end =   [1,2,3, 4,5,6, 7,8,9,10, 14,15,16, 11,12,13]
    ske_predicted = torch.zeros(len(start))
    ske_target = torch.zeros(len(start))
    for i in range(len(start)):
        ske_predicted[i] = torch.mean(torch.linalg.norm(predicted[:,:,start[i],:] - predicted[:,:,end[i],:], dim=2)).contiguous()
        ske_target[i] = torch.mean(torch.linalg.norm(target[:,:,start[i],:] - target[:,:,end[i],:], dim=2)).contiguous()



    return torch.mean(torch.linalg.norm(ske_predicted[i] - ske_target[i]))
//  [python]
def gradient_forward_(self, out, grad_output, out_ForwPass, noise_dir=None, noise_eps=None):
        gradients_inverse = []
        for layer in self.layers:
            out, grad_output, layergrad = layer.gradient_forward_(out, grad_output)
            gradients_inverse += layergrad
        gradients_forward = []
        # re-use output of forward pass (from computation of loss) here to
        # reduce the numerical errors during inverse pass
        # (inverse pass is done during backprop to recomputate activations)
        out = out_ForwPass
        out_test = out - noise_eps * noise_dir # remove influence of noise
        recon_error = (torch.norm(out_ForwPass - out_test) / out_ForwPass.shape[0]).item()
        print('Recon error in inverse: ' + str(recon_error))
        for layers in reversed(self.layers):
            out, grad_output, layergrad = layers.gradient_(out, grad_output)
            gradients_forward += layergrad
        gradients_forward = [g for g in gradients_forward if g is not None]
        gradients_inverse = [g for g in gradients_inverse if g is not None]
        gradients = [g_inv + g_forw for (g_inv, g_forw) in zip(gradients_inverse, reversed(gradients_forward))]
        return grad_output, gradients, out
//  [python]
def gradient_forward_(self, out, grad_output, out_ForwPass, noise_dir=None, noise_eps=None):
        gradients_inverse = []
        for layer in self.layers:
            out, grad_output, layergrad = layer.gradient_forward_(out, grad_output)
            gradients_inverse += layergrad
        gradients_forward = []
        # re-use output of forward pass (from computation of loss) here to
        # reduce the numerical errors during inverse pass
        # (inverse pass is done during backprop to recomputate activations)
        out = out_ForwPass
        out_test = out - noise_eps * noise_dir # remove influence of noise
        recon_error = (torch.linalg.norm(out_ForwPass - out_test) / out_ForwPass.shape[0]).item()
        print('Recon error in inverse: ' + str(recon_error))
        for layers in reversed(self.layers):
            out, grad_output, layergrad = layers.gradient_(out, grad_output)
            gradients_forward += layergrad
        gradients_forward = [g for g in gradients_forward if g is not None]
        gradients_inverse = [g for g in gradients_inverse if g is not None]
        gradients = [g_inv + g_forw for (g_inv, g_forw) in zip(gradients_inverse, reversed(gradients_forward))]
        return grad_output, gradients, out
//  [python]
def mpjpe_by_action(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    dist = torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 2)

    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name].update(torch.mean(dist).item()*num, num)
    else:
        for i in range(num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name].update(dist[i].item(), 1)

    return action_error_sum
//  [python]
def mpjpe_by_action(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    dist = torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 2)

    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name].update(torch.mean(dist).item()*num, num)
    else:
        for i in range(num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name].update(dist[i].item(), 1)

    return action_error_sum
//  [python]
def get_morphed_w_code(self, new_w_code, fixed_w):
        interpolation_direction = new_w_code - fixed_w
        interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
        direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
        result_w = fixed_w + direction_to_move
        self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code

        return result_w
//  [python]
def get_morphed_w_code(self, new_w_code, fixed_w):
        interpolation_direction = new_w_code - fixed_w
        interpolation_direction_norm = torch.linalg.norm(interpolation_direction, p=2)
        direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
        result_w = fixed_w + direction_to_move
        self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code

        return result_w
//  [python]
def mpjpe_by_joint_mae(predicted, target,num):
    assert predicted.shape == target.shape
    # this is the joint
    mpjpe_joint = torch.mean(torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 3),dim=len(target.shape)-4)
    print("\nthe mpjpe/joint",mpjpe_joint)
    # this is the order of joint from big to small
    index = torch.flip(mpjpe_joint.sort(-1).indices,dims=[0])
    index = index.split(num,-1)[0]
    print("\nerror joint",index)
    return index
//  [python]
def mpjpe_by_joint_mae(predicted, target,num):
    assert predicted.shape == target.shape
    # this is the joint
    mpjpe_joint = torch.mean(torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 3),dim=len(target.shape)-4)
    print("\nthe mpjpe/joint",mpjpe_joint)
    # this is the order of joint from big to small
    index = torch.flip(mpjpe_joint.sort(-1).indices,dims=[0])
    index = index.split(num,-1)[0]
    print("\nerror joint",index)
    return index
//  [python]
def RBF(self, positions):
        positions = positions.reshape(self.n_replicas, self.n_particles * 3)
        samples1 = torch.unsqueeze(positions, 0)
        samples2 = torch.unsqueeze(positions, 1)
        return torch.exp(-self.gamma*torch.norm(samples1 - samples2, p=2, dim=2) ** 2)
//  [python]
def RBF(self, positions):
        positions = positions.reshape(self.n_replicas, self.n_particles * 3)
        samples1 = torch.unsqueeze(positions, 0)
        samples2 = torch.unsqueeze(positions, 1)
        return torch.exp(-self.gamma*torch.linalg.norm(samples1 - samples2, p=2, dim=2) ** 2)
//  [python]
def _domain_difference_img(self, original, transformed, bandwidth=3):
        """
            Compute difference image in one domain between original image
            in that domain and the transformed image from the other domain.
            Bandwidth governs the norm difference clipping threshold
        """
        d = torch.norm(original - transformed, p=2, dim=1)
        threshold = torch.mean(d) + bandwidth * torch.std(d)
        d = torch.where(d < threshold, d, threshold)
        return d / torch.max(d)
//  [python]
def _domain_difference_img(self, original, transformed, bandwidth=3):
        """
            Compute difference image in one domain between original image
            in that domain and the transformed image from the other domain.
            Bandwidth governs the norm difference clipping threshold
        """
        d = torch.linalg.norm(original - transformed, p=2, dim=1)
        threshold = torch.mean(d) + bandwidth * torch.std(d)
        d = torch.where(d < threshold, d, threshold)
        return d / torch.max(d)
//  [python]
def norm(x, p='fro', axis=None, keepdims=False):
  """
  x : must be 1-D or 2-D
  """
  if torch.is_tensor(x):
    return torch.norm(x, p=p, dim=axis, keepdim=keepdims)
  if isinstance(axis, (tuple, list)) and len(axis) == x.ndim:
    axis = None
  if p == 'fro':
    p = 'euclidean'
  if tf.is_tensor(x):
    return tf.norm(x, ord=p, axis=axis, keepdims=keepdims)
  return tf.norm(x, ord=p, axis=axis, keepdims=keepdims).numpy()
//  [python]
def norm(x, p='fro', axis=None, keepdims=False):
  """
  x : must be 1-D or 2-D
  """
  if torch.is_tensor(x):
    return torch.linalg.norm(x, ord=p, dim=axis, keepdim=keepdims)
  if isinstance(axis, (tuple, list)) and len(axis) == x.ndim:
    axis = None
  if p == 'fro':
    p = 'euclidean'
  if tf.is_tensor(x):
    return tf.norm(x, ord=p, axis=axis, keepdims=keepdims)
  return tf.norm(x, ord=p, axis=axis, keepdims=keepdims).numpy()
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    assert children[0] == 3
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, int(children[0])].clone()
    spine_rest_loc = rel_rest_pose[:, int(children[0])].clone()
    spine_norm = torch.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine)
    center_final_loc = 0
    center_rest_loc = 0
    for child in pelvis_child:
        if child == int(children[0]):
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()
    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.norm(center_rest_loc, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    assert children[0] == 3
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, int(children[0])].clone()
    spine_rest_loc = rel_rest_pose[:, int(children[0])].clone()
    spine_norm = torch.linalg.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine)
    center_final_loc = 0
    center_rest_loc = 0
    for child in pelvis_child:
        if child == int(children[0]):
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()
    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.linalg.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.linalg.norm(center_rest_loc, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def xyzquat_to_tf(position_quat: torch.Tensor) -> torch.Tensor:
    """
    convert [x, y, z, qw, qx, qy, qz] to 4 x 4 transformation matrices
    """
    try:
        position_quat[:, 3:] = position_quat[:, 3:] / torch.norm(
            position_quat[:, 3:], dim=1
        ).reshape(-1, 1)
        th_T = th.SE3(x_y_z_quaternion=position_quat).to_matrix()
    except ValueError:
        print("Zero quat error!")
    return th_T.squeeze()
//  [python]
def xyzquat_to_tf(position_quat: torch.Tensor) -> torch.Tensor:
    """
    convert [x, y, z, qw, qx, qy, qz] to 4 x 4 transformation matrices
    """
    try:
        position_quat[:, 3:] = position_quat[:, 3:] / torch.linalg.norm(
            position_quat[:, 3:], dim=1
        ).reshape(-1, 1)
        th_T = th.SE3(x_y_z_quaternion=position_quat).to_matrix()
    except ValueError:
        print("Zero quat error!")
    return th_T.squeeze()
//  [python]
def delta_measure(self, C, gamma, epsilon=0.01):
        """
        Delta measure (sensibility of Frobenius norm to change of aperture)
        :param C: Conceptor object
        :param gamma: Gamma
        :param epsilon: Epsilon
        :return: Delta measure
        """
        # Conceptors with two gammas
        A = Conceptor.operator_PHI(C, gamma - epsilon)
        B = Conceptor.operator_PHI(C, gamma + epsilon)

        # Gradient of Frobenius norm
        A_norm = math.pow(torch.norm(A, p=2), 2)
        B_norm = math.pow(torch.norm(B, p=2), 2)
        d_C_norm = B_norm - A_norm

        # Change in log (gamma)
        d_log_gamma = torch.log(gamma + epsilon) - torch.log(gamma - epsilon)
        return d_C_norm / d_log_gamma, d_C_norm
//  [python]
def delta_measure(self, C, gamma, epsilon=0.01):
        """
        Delta measure (sensibility of Frobenius norm to change of aperture)
        :param C: Conceptor object
        :param gamma: Gamma
        :param epsilon: Epsilon
        :return: Delta measure
        """
        # Conceptors with two gammas
        A = Conceptor.operator_PHI(C, gamma - epsilon)
        B = Conceptor.operator_PHI(C, gamma + epsilon)

        # Gradient of Frobenius norm
        A_norm = math.pow(torch.linalg.norm(A, p=2), 2)
        B_norm = math.pow(torch.linalg.norm(B, p=2), 2)
        d_C_norm = B_norm - A_norm

        # Change in log (gamma)
        d_log_gamma = torch.log(gamma + epsilon) - torch.log(gamma - epsilon)
        return d_C_norm / d_log_gamma, d_C_norm
//  [python]
def getInvariant(self, positions):
        new_pos = torch.zeros(self.n_replicas, 1, self.n_particles*3)
        for i in range(len(positions)):
            inp = positions[i].view(-1, self.n_particles, 3)-self.offset
            com = torch.mean(inp, dim=1)
            centered = inp-com[:,None,:]
            
            top = centered[:,-1]
            A = top / torch.norm(top, dim=1)[:, None]

            B = torch.tensor([1., 0., 0.], device=A.device)

            B = B.type(A.dtype)

            dot = (A @ B)
            cross = torch.cross(A, B.repeat(A.size(0), 1))
            cross_norm = torch.norm(cross, dim=1)

            zeros = torch.zeros_like(dot, device=dot.device)
            ones = torch.ones_like(dot, device=dot.device)
            G1 = torch.vstack((dot, -cross_norm, zeros)).T
            G2 = torch.vstack((cross_norm, dot, zeros)).T
            G3 = torch.vstack((zeros, zeros, ones)).T
            G = torch.stack((G1, G2, G3), dim=2)

            u = A
            v = (B - (dot[:, None] * A)) / torch.norm((B - (dot[:, None] * A)), dim=1)[:, None]
            w = torch.cross(B.repeat(A.size(0), 1), A)
            F_inv = torch.stack([u, v, w], dim=2)
            F = torch.inverse(F_inv)
            U = torch.matmul(torch.matmul(F_inv, G), F)
            rotated = torch.matmul(centered, U)
            new_pos[i] = rotated.view((1, self.n_particles*3))
        return new_pos
//  [python]
def getInvariant(self, positions):
        new_pos = torch.zeros(self.n_replicas, 1, self.n_particles*3)
        for i in range(len(positions)):
            inp = positions[i].view(-1, self.n_particles, 3)-self.offset
            com = torch.mean(inp, dim=1)
            centered = inp-com[:,None,:]
            
            top = centered[:,-1]
            A = top / torch.linalg.norm(top, dim=1)[:, None]

            B = torch.tensor([1., 0., 0.], device=A.device)

            B = B.type(A.dtype)

            dot = (A @ B)
            cross = torch.cross(A, B.repeat(A.size(0), 1))
            cross_norm = torch.linalg.norm(cross, dim=1)

            zeros = torch.zeros_like(dot, device=dot.device)
            ones = torch.ones_like(dot, device=dot.device)
            G1 = torch.vstack((dot, -cross_norm, zeros)).T
            G2 = torch.vstack((cross_norm, dot, zeros)).T
            G3 = torch.vstack((zeros, zeros, ones)).T
            G = torch.stack((G1, G2, G3), dim=2)

            u = A
            v = (B - (dot[:, None] * A)) / torch.linalg.norm((B - (dot[:, None] * A)), dim=1)[:, None]
            w = torch.cross(B.repeat(A.size(0), 1), A)
            F_inv = torch.stack([u, v, w], dim=2)
            F = torch.inverse(F_inv)
            U = torch.matmul(torch.matmul(F_inv, G), F)
            rotated = torch.matmul(centered, U)
            new_pos[i] = rotated.view((1, self.n_particles*3))
        return new_pos
//  [python]
u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)
u_exact = torch.from_numpy(u_exact)
error = torch.abs(u_pred - u_exact) # get the error in every points
error_t = torch.norm(error)/torch.norm(u_exact)
//  [python]
u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)
u_exact = torch.from_numpy(u_exact)
error = torch.abs(u_pred - u_exact) # get the error in every points
error_t = torch.linalg.norm(error)/torch.linalg.norm(u_exact)
//  [python]
def clean_up_optitrack(poses):
    """
    Filter large jumps in mocap data
    """
    traj_sz = poses.shape[0]
    diff_pose_mags = []
    adjusted_count = 0
    filtered_poses = torch.empty((0, 4, 4), device=poses.device)
    idx_removed = []
    for i in range(traj_sz):
        if i > 0:
            diff_pose = torch.inverse(poses[i - 1, :]) @ poses[i, :]
            diff_pose_mag = torch.norm(diff_pose[:3, 3])
            diff_pose_mags.append(diff_pose_mag)
            avg_diff_pose_mag = sum(diff_pose_mags) / len(diff_pose_mags)
            if i > 1 and diff_pose_mag > 10 * avg_diff_pose_mag:
                adjusted_count += 1
                idx_removed.append(i)
            else:
                filtered_poses = torch.cat(
                    (filtered_poses, poses[i, :][None, :]), dim=0
                )
                # print(f"Jump @ t = {i} : avg: {avg_diff_pose_mag}, curr: {diff_pose_mag}")
                # poses[i, :] = tf_to_xyzquat(xyzquat_to_tf(poses[i - 1, :]) @ xyzquat_to_tf(prev_diff_pose))
            # prev_diff_pose = diff_pose
    print(f"Adjusted {adjusted_count} / {traj_sz} object-sensor poses")
    return filtered_poses, idx_removed
//  [python]
def clean_up_optitrack(poses):
    """
    Filter large jumps in mocap data
    """
    traj_sz = poses.shape[0]
    diff_pose_mags = []
    adjusted_count = 0
    filtered_poses = torch.empty((0, 4, 4), device=poses.device)
    idx_removed = []
    for i in range(traj_sz):
        if i > 0:
            diff_pose = torch.inverse(poses[i - 1, :]) @ poses[i, :]
            diff_pose_mag = torch.linalg.norm(diff_pose[:3, 3])
            diff_pose_mags.append(diff_pose_mag)
            avg_diff_pose_mag = sum(diff_pose_mags) / len(diff_pose_mags)
            if i > 1 and diff_pose_mag > 10 * avg_diff_pose_mag:
                adjusted_count += 1
                idx_removed.append(i)
            else:
                filtered_poses = torch.cat(
                    (filtered_poses, poses[i, :][None, :]), dim=0
                )
                # print(f"Jump @ t = {i} : avg: {avg_diff_pose_mag}, curr: {diff_pose_mag}")
                # poses[i, :] = tf_to_xyzquat(xyzquat_to_tf(poses[i - 1, :]) @ xyzquat_to_tf(prev_diff_pose))
            # prev_diff_pose = diff_pose
    print(f"Adjusted {adjusted_count} / {traj_sz} object-sensor poses")
    return filtered_poses, idx_removed
//  [python]
def test_jacobian_pullback_dense():
    for get_task in linear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        pull_back = PullBackDense(generator=generator, examples=loader)
        push_forward = PushForwardDense(generator=generator, examples=loader)
        dw = random_pvector(lc, device=device)

        doutput_lin = push_forward.mv(dw)
        dinput_lin = pull_back.mv(doutput_lin)
        check_ratio(
            torch.dot(
                dw.get_flat_representation(), dinput_lin.get_flat_representation()
            ),
            torch.norm(doutput_lin.get_flat_representation()) ** 2,
        )
//  [python]
def test_jacobian_pullback_dense():
    for get_task in linear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        pull_back = PullBackDense(generator=generator, examples=loader)
        push_forward = PushForwardDense(generator=generator, examples=loader)
        dw = random_pvector(lc, device=device)

        doutput_lin = push_forward.mv(dw)
        dinput_lin = pull_back.mv(doutput_lin)
        check_ratio(
            torch.dot(
                dw.get_flat_representation(), dinput_lin.get_flat_representation()
            ),
            torch.linalg.norm(doutput_lin.get_flat_representation()) ** 2,
        )
//  [python]
def evaluate():
    N_test = 100
    x = np.linspace(-1, 1, N_test)
    y = np.linspace(-1, 1, N_test)
    x, y = np.meshgrid(x, y)
    xy_test = np.stack((x.flatten(), y.flatten()),1)
    xy_test = torch.from_numpy(xy_test).cuda() # to the model for the prediction
    
    u_pred = pred(xy_test)
    u_pred = u_pred.data
    u_pred = u_pred.reshape(N_test, N_test).cpu()
    

    u_exact = np.zeros(x.shape)
    u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
    u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)

    u_exact = u_exact.reshape(x.shape)
    u_exact = torch.from_numpy(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.norm(error)/torch.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def evaluate():
    N_test = 100
    x = np.linspace(-1, 1, N_test)
    y = np.linspace(-1, 1, N_test)
    x, y = np.meshgrid(x, y)
    xy_test = np.stack((x.flatten(), y.flatten()),1)
    xy_test = torch.from_numpy(xy_test).cuda() # to the model for the prediction
    
    u_pred = pred(xy_test)
    u_pred = u_pred.data
    u_pred = u_pred.reshape(N_test, N_test).cpu()
    

    u_exact = np.zeros(x.shape)
    u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
    u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)

    u_exact = u_exact.reshape(x.shape)
    u_exact = torch.from_numpy(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.linalg.norm(error)/torch.linalg.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def test_jacobian_pblockdiag_vs_pdense():
    for get_task in linear_tasks + nonlinear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_blockdiag = PMatBlockDiag(generator=generator, examples=loader)
        PMat_dense = PMatDense(generator=generator, examples=loader)

        # Test get_dense_tensor
        matrix_blockdiag = PMat_blockdiag.get_dense_tensor()
        matrix_dense = PMat_dense.get_dense_tensor()
        for layer_id, layer in lc.layers.items():
            start = lc.p_pos[layer_id]
            # compare blocks
            check_tensors(
                matrix_dense[
                    start : start + layer.numel(), start : start + layer.numel()
                ],
                matrix_blockdiag[
                    start : start + layer.numel(), start : start + layer.numel()
                ],
            )
            # verify that the rest is 0
            assert (
                torch.norm(
                    matrix_blockdiag[
                        start : start + layer.numel(), start + layer.numel() :
                    ]
                )
                < 1e-5
            )
            assert (
                torch.norm(
                    matrix_blockdiag[
                        start + layer.numel() :, start : start + layer.numel()
                    ]
                )
                < 1e-5
            )
//  [python]
def test_jacobian_pblockdiag_vs_pdense():
    for get_task in linear_tasks + nonlinear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_blockdiag = PMatBlockDiag(generator=generator, examples=loader)
        PMat_dense = PMatDense(generator=generator, examples=loader)

        # Test get_dense_tensor
        matrix_blockdiag = PMat_blockdiag.get_dense_tensor()
        matrix_dense = PMat_dense.get_dense_tensor()
        for layer_id, layer in lc.layers.items():
            start = lc.p_pos[layer_id]
            # compare blocks
            check_tensors(
                matrix_dense[
                    start : start + layer.numel(), start : start + layer.numel()
                ],
                matrix_blockdiag[
                    start : start + layer.numel(), start : start + layer.numel()
                ],
            )
            # verify that the rest is 0
            assert (
                torch.linalg.norm(
                    matrix_blockdiag[
                        start : start + layer.numel(), start + layer.numel() :
                    ]
                )
                < 1e-5
            )
            assert (
                torch.linalg.norm(
                    matrix_blockdiag[
                        start + layer.numel() :, start : start + layer.numel()
                    ]
                )
                < 1e-5
            )
//  [python]
def _sample_orthonormal_to(self, mu, dim):
        """Sample point on sphere orthogonal to mu.
        """
        v = GVar(torch.randn(dim))
        rescale_value = mu.dot(v) / mu.norm()
        proj_mu_v = mu * rescale_value.expand(dim)
        ortho = v - proj_mu_v
        ortho_norm = torch.norm(ortho)
        return ortho / ortho_norm.expand_as(ortho)
//  [python]
def _sample_orthonormal_to(self, mu, dim):
        """Sample point on sphere orthogonal to mu.
        """
        v = GVar(torch.randn(dim))
        rescale_value = mu.dot(v) / torch.linalg.norm(mu)
        proj_mu_v = mu * rescale_value.expand(dim)
        ortho = v - proj_mu_v
        ortho_norm = torch.linalg.norm(ortho)
        return ortho / ortho_norm.expand_as(ortho)
//  [python]
def test_jacobian_pquasidiag():
    for get_task in [get_conv_task, get_fullyconnect_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_qd = PMatQuasiDiag(generator=generator, examples=loader)
        dense_tensor = PMat_qd.get_dense_tensor()

        v = random_pvector(lc, device=device)
        v_flat = v.get_flat_representation()

        check_tensors(torch.diag(dense_tensor), PMat_qd.get_diag())

        check_ratio(torch.norm(dense_tensor), PMat_qd.frobenius_norm())

        check_ratio(torch.trace(dense_tensor), PMat_qd.trace())

        mv = PMat_qd.mv(v)
        check_tensors(torch.mv(dense_tensor, v_flat), mv.get_flat_representation())

        check_ratio(torch.dot(torch.mv(dense_tensor, v_flat), v_flat), PMat_qd.vTMv(v))

        # Test solve
        regul = 1e-8
        v_back = PMat_qd.solve(mv + regul * v, regul=regul)
        check_tensors(v.get_flat_representation(), v_back.get_flat_representation())
//  [python]
def test_jacobian_pquasidiag():
    for get_task in [get_conv_task, get_fullyconnect_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_qd = PMatQuasiDiag(generator=generator, examples=loader)
        dense_tensor = PMat_qd.get_dense_tensor()

        v = random_pvector(lc, device=device)
        v_flat = v.get_flat_representation()

        check_tensors(torch.diag(dense_tensor), PMat_qd.get_diag())

        check_ratio(torch.linalg.norm(dense_tensor), PMat_qd.frobenius_norm())

        check_ratio(torch.trace(dense_tensor), PMat_qd.trace())

        mv = PMat_qd.mv(v)
        check_tensors(torch.mv(dense_tensor, v_flat), mv.get_flat_representation())

        check_ratio(torch.dot(torch.mv(dense_tensor, v_flat), v_flat), PMat_qd.vTMv(v))

        # Test solve
        regul = 1e-8
        v_back = PMat_qd.solve(mv + regul * v, regul=regul)
        check_tensors(v.get_flat_representation(), v_back.get_flat_representation())
//  [python]
def exp_SO3_b(phi):
    eps = 1e-8
    C = torch.zeros(phi.size(0), 3, 3, device=phi.device)

    phi_norm = torch.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    phi_no_sel = phi[~sel]

    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        unit_phi_skewed_sel = skew3_b(unit_phi_sel)
        C[sel] = torch.eye(3, 3, device=phi.device).repeat([phi_norm_sel.size(0), 1, 1]) + \
                 torch.sin(phi_norm_sel) * unit_phi_skewed_sel + \
                 (1 - torch.cos(phi_norm_sel)) * torch.matmul(unit_phi_skewed_sel, unit_phi_skewed_sel)

    if phi_no_sel.size(0):
        phi_skewed_no_sel = skew3_b(phi_no_sel)
        C[~sel] = torch.eye(3, 3, device=phi.device).repeat([phi_no_sel.size(0), 1, 1]) + phi_skewed_no_sel

    return C
//  [python]
def exp_SO3_b(phi):
    eps = 1e-8
    C = torch.zeros(phi.size(0), 3, 3, device=phi.device)

    phi_norm = torch.linalg.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    phi_no_sel = phi[~sel]

    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        unit_phi_skewed_sel = skew3_b(unit_phi_sel)
        C[sel] = torch.eye(3, 3, device=phi.device).repeat([phi_norm_sel.size(0), 1, 1]) + \
                 torch.sin(phi_norm_sel) * unit_phi_skewed_sel + \
                 (1 - torch.cos(phi_norm_sel)) * torch.matmul(unit_phi_skewed_sel, unit_phi_skewed_sel)

    if phi_no_sel.size(0):
        phi_skewed_no_sel = skew3_b(phi_no_sel)
        C[~sel] = torch.eye(3, 3, device=phi.device).repeat([phi_no_sel.size(0), 1, 1]) + phi_skewed_no_sel

    return C
//  [python]
def vmf_unif_sampler(self, mu):
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = mu[i].norm().expand(id_dim)
            munoise = self.add_norm_noise(munorm, self.norm_eps)
            if float(mu[i].norm().data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                w = self._sample_weight(self.kappa, id_dim)
                wtorch = GVar(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)

                # compute new point
                scale_factr = torch.sqrt(GVar(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr
                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munoise
            else:
                rand_draw = GVar(torch.randn(id_dim))
                rand_draw = rand_draw / torch.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * GVar(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0)
//  [python]
def vmf_unif_sampler(self, mu):
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = torch.linalg.norm(mu[i]).expand(id_dim)
            munoise = self.add_norm_noise(munorm, self.norm_eps)
            if float(torch.linalg.norm(mu[i]).data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                w = self._sample_weight(self.kappa, id_dim)
                wtorch = GVar(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)

                # compute new point
                scale_factr = torch.sqrt(GVar(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr
                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munoise
            else:
                rand_draw = GVar(torch.randn(id_dim))
                rand_draw = rand_draw / torch.linalg.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * GVar(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0)
//  [python]
def length(a, b):
    return torch.norm(a-b, dim=-1)
//  [python]
def length(a, b):
    return torch.linalg.norm(a-b, dim=-1)
//  [python]
def vmf_sampler(self, mu):
        mu = mu.cpu()
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = mu[i].norm().expand(id_dim)  # TODO norm p=?
            if float(mu[i].norm().data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                # w = self._sample_weight(self.kappa, id_dim)     # TODO mine?

                w = vMF.sample_vmf_w(self.kappa, id_dim)

                wtorch = GVar(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)
                # v= vMF.sample_vmf_v(mu[i])

                # compute new point
                scale_factr = torch.sqrt(GVar(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr

                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munorm
            else:
                rand_draw = GVar(torch.randn(id_dim))
                rand_draw = rand_draw / torch.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * GVar(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0).cuda()
//  [python]
def vmf_sampler(self, mu):
        mu = mu.cpu()
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = torch.linalg.norm(mu[i]).expand(id_dim)  # TODO norm p=?
            if float(torch.linalg.norm(mu[i]).data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                # w = self._sample_weight(self.kappa, id_dim)     # TODO mine?

                w = vMF.sample_vmf_w(self.kappa, id_dim)

                wtorch = GVar(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)
                # v= vMF.sample_vmf_v(mu[i])

                # compute new point
                scale_factr = torch.sqrt(GVar(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr

                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munorm
            else:
                rand_draw = GVar(torch.randn(id_dim))
                rand_draw = rand_draw / torch.linalg.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * GVar(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0).cuda()