return (x1[1] * x2[0] - x1[0] * x2[1]) / (np.square(x1[0]) +
np.square(x2[0]))
register_diff(np.add, lambda x1, x2: x1[1] + x2[1])
register_diff(np.subtract, lambda x1, x2: x1[1] - x2[1])
register_diff(np.multiply, multiply_diff)
register_diff(np.matmul, matmul_diff)
register_diff(np.divide, divide_diff)
register_diff(np.true_divide, divide_diff)
register_diff(np.power, pow_diff)
register_diff(np.positive, lambda x: +x[1])
register_diff(np.negative, lambda x: -x[1])
register_diff(np.conj, lambda x: np.conj(x[1]))
register_diff(np.conj, lambda x: np.conj(x[1]))
register_diff(np.exp, lambda x: x[1] * np.exp(x[0]))
register_diff(np.exp2, lambda x: x[1] * np.log(2) * np.exp(x[0]))
register_diff(np.log, lambda x: x[1] / x[0])
register_diff(np.log2, lambda x: x[1] / (np.log(2) * x[0]))
register_diff(np.log10, lambda x: x[1] / (np.log(10) * x[0]))
register_diff(np.sqrt, lambda x: x[1] / (2 * np.sqrt(x[0])))
register_diff(np.square, lambda x: 2 * x[1] * x[0])
register_diff(np.cbrt, lambda x: x[1] / (3 * (x[0]**(2 / 3))))
register_diff(np.reciprocal, lambda x: -x[1] / np.square(x[0]))
register_diff(np.broadcast_to, lambda x, shape: np.broadcast_to(x[1], shape))
register_diff(np.sin, lambda x: x[1] * np.cos(x[0]))
register_diff(np.cos, lambda x: -x[1] * np.sin(x[0]))
register_diff(np.tan, lambda x: x[1] / np.square(np.cos(x[0])))
register_diff(np.arcsin, lambda x: x[1] / np.sqrt(1 - np.square(x[0])))
register_diff(np.arccos, lambda x: -x[1] / np.sqrt(1 - np.square(x[0])))
lambda ans, x, y: lambda g: g * balanced_eq(x, ans, y),
lambda ans, x, y: lambda g: g * balanced_eq(y, ans, x),
)
defjvp(
np.fmax,
lambda ans, x, y: lambda g: g * balanced_eq(x, ans, y),
lambda ans, x, y: lambda g: g * balanced_eq(y, ans, x),
)
defjvp(
np.fmin,
lambda ans, x, y: lambda g: g * balanced_eq(x, ans, y),
lambda ans, x, y: lambda g: g * balanced_eq(y, ans, x),
)
defjvp(
np.logaddexp,
lambda ans, x, y: lambda g: g * np.exp(x - ans),
lambda ans, x, y: lambda g: g * np.exp(y - ans),
)
defjvp(
np.logaddexp2,
lambda ans, x, y: lambda g: g * 2**(x - ans),
lambda ans, x, y: lambda g: g * 2**(y - ans),
)
defjvp(np.true_divide, "same", lambda ans, x, y: lambda g: -g * x / y**2)
defjvp(
np.mod,
lambda ans, x, y: lambda g: np.broadcast_to(g, np.shape(ans)),
lambda ans, x, y: lambda g: -g * np.floor(x / y),
)
defjvp(
np.remainder,
lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),
lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),
)
defvjp(
np.fmax,
lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),
lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),
)
defvjp(
np.fmin,
lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),
lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),
)
defvjp(
np.logaddexp,
lambda ans, x, y: unbroadcast_f(x, lambda g: g * np.exp(x - ans)),
lambda ans, x, y: unbroadcast_f(y, lambda g: g * np.exp(y - ans)),
)
defvjp(
np.logaddexp2,
lambda ans, x, y: unbroadcast_f(x, lambda g: g * 2 ** (x - ans)),
lambda ans, x, y: unbroadcast_f(y, lambda g: g * 2 ** (y - ans)),
)
defvjp(
np.true_divide,
lambda ans, x, y: unbroadcast_f(x, lambda g: g / y),
lambda ans, x, y: unbroadcast_f(y, lambda g: -g * x / y ** 2),
)
defvjp(
np.mod,
lambda ans, x, y: unbroadcast_f(x, lambda g: g),
)
register_diff(np.add, lambda x1, x2: x1[1] + x2[1])
register_diff(np.subtract, lambda x1, x2: x1[1] - x2[1])
register_diff(np.multiply, multiply_diff)
register_diff(np.matmul, matmul_diff)
register_diff(np.divide, divide_diff)
register_diff(np.true_divide, divide_diff)
register_diff(np.power, pow_diff)
register_diff(
np.absolute,
lambda x: x[1] * np.where(np.sign(x[0]) == 0, float("nan"), np.sign(x[0])),
)
register_diff(np.positive, lambda x: +x[1])
register_diff(np.negative, lambda x: -x[1])
register_diff(np.conj, lambda x: np.conj(x[1]))
register_diff(np.exp, lambda x: x[1] * np.exp(x[0]))
register_diff(np.exp2, lambda x: x[1] * np.log(2) * np.exp2(x[0]))
register_diff(np.log, lambda x: x[1] / x[0])
register_diff(np.log2, lambda x: x[1] / (np.log(2) * x[0]))
register_diff(np.log10, lambda x: x[1] / (np.log(10) * x[0]))
register_diff(np.sqrt, lambda x: x[1] / (2 * np.sqrt(x[0])))
register_diff(np.square, lambda x: 2 * x[1] * x[0])
register_diff(np.cbrt, lambda x: x[1] / (3 * (x[0]**(2 / 3))))
register_diff(np.reciprocal, lambda x: -x[1] / np.square(x[0]))
register_diff(np.broadcast_to, lambda x, shape: np.broadcast_to(x[1], shape))
register_diff(np.sin, lambda x: x[1] * np.cos(x[0]))
register_diff(np.cos, lambda x: -x[1] * np.sin(x[0]))
register_diff(np.tan, lambda x: x[1] / np.square(np.cos(x[0])))
register_diff(np.arcsin, lambda x: x[1] / np.sqrt(1 - np.square(x[0])))
register_diff(np.arccos, lambda x: -x[1] / np.sqrt(1 - np.square(x[0])))