def erf(F, x: Tensor): if MXNET_HAS_ERF: return F.erf(x) # Using numerical recipes approximation for erf function # accurate to 1E-7 ones = x.ones_like() zeros = x.zeros_like() t = ones / (ones + 0.5 * x.abs()) coefficients = [ 1.00002368, 0.37409196, 0.09678418, -0.18628806, 0.27886807, -1.13520398, 1.48851587, -0.82215223, 0.17087277, ] inner = zeros for c in coefficients[::-1]: inner = t * (c + inner) res = ones - t * (inner - 1.26551223 - x.square()).exp() return F.where(F.broadcast_greater_equal(x, zeros), res, -1.0 * res)
def mahalanobis_distance( F, W: Tensor, D: Tensor, capacitance_tril: Tensor, x: Tensor ) -> Tensor: r""" Uses the Woodbury matrix identity .. math:: (W W^T + D)^{-1} = D^{-1} - D^{-1} W C^{-1} W^T D^{-1}, where :math:`C` is the capacitance matrix :math:`I + W^T D^{-1} W`, to compute the squared Mahalanobis distance :math:`x^T (W W^T + D)^{-1} x`. Parameters ---------- F W (..., dim, rank) D (..., dim) capacitance_tril (..., rank, rank) x (..., dim) Returns ------- """ xx = x.expand_dims(axis=-1) # (..., rank, 1) Wt_Dinv_x = F.linalg_gemm2( F.broadcast_div(W, D.expand_dims(axis=-1)), xx, transpose_a=True ) # compute x^T D^-1 x, (...,) maholanobis_D_inv = F.broadcast_div(x.square(), D).sum(axis=-1) # (..., rank) L_inv_Wt_Dinv_x = F.linalg_trsm(capacitance_tril, Wt_Dinv_x).squeeze( axis=-1 ) maholanobis_L = L_inv_Wt_Dinv_x.square().sum(axis=-1).squeeze() return F.broadcast_minus(maholanobis_D_inv, maholanobis_L)