def loss_fun(delta: ep.Tensor, logits: ep.Tensor) -> ep.Tensor:
assert x.shape[0] == logits.shape[0]
assert delta.shape == x.shape
x_hat = x + delta
logits_hat = model(x_hat)
loss = ep.kl_div_with_logits(logits, logits_hat).sum()
return loss
def loss_fn(inputs):
logits_q = model(inputs)
return ep.kl_div_with_logits(logits_p, logits_q).sum()
def test_kl_div_with_logits(dummy: ep.Tensor, axis: int) -> None:
logits_p = logits_q = ep.arange(dummy, 12).float32().reshape((3, 4))
assert (ep.kl_div_with_logits(logits_p, logits_q, axis=axis) == 0).all()
