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)
