Make kl_div accept target in log space

[nikitaved]

Fixes 32520, implements 34536.

Here are some benchmarks:

import torch
import torch.nn.functional as F
from IPython import get_ipython

ipython = get_ipython()

torch.set_num_threads(1)

for d in [5, 10, 20, 50, 100, 1000]:
    i = torch.rand(d, d)
    t = torch.rand(d, d)
    print(f"Size: {d}x{d}")
    ipython.magic("timeit F.kl_div(i, t, reduction='none', log_target=False)")
    ipython.magic("timeit F.kl_div(i, t.log(), reduction='none', log_target=True)")

Output:

Size: 5x5
16 µs ± 33 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
8.24 µs ± 17.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Size: 10x10
16.7 µs ± 17.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
8.7 µs ± 20.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Size: 20x20
17.7 µs ± 47.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
9.7 µs ± 28.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Size: 50x50
23.6 µs ± 60.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
15 µs ± 33.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Size: 100x100
42.8 µs ± 223 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
34 µs ± 17.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Size: 1000x1000
3.9 ms ± 1.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.45 ms ± 364 ns per loop (mean ± std. dev. of 7 runs, 100 loops each)

[dr-ci]

💊 CircleCI build failures summary and remediations

As of commit cf4814f (more details on the Dr. CI page):

---

• 1/1 failures introduced in this PR

---

XLA failure

Job pytorch_xla_linux_xenial_py3_6_clang7_build is failing. Please create an issue with title prefixed by [PT_BREAK] in pytorch/xla and link to to this PR. If you have questions, please reach out to @ailzhang / @dlibenzi / @JackCaoG.

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker.

See how this bot performed.

This comment has been revised 129 times.

[nikitaved]

@ezyang, @ngimel, could you please tell what would be the best way to handle buffers in this case? I need to store that boolean flag for backward.

[ezyang]

cc @ailzhang @dlibenzi this is going to affect xla bindings

[ezyang]

Needs some benchmarks for the new codepath. Compare log_space=False/True

[ezyang]

tests, benchmarks

[nikitaved]

Needs some benchmarks for the new codepath. Compare log_space=False/True

If I sample from Dirichlet distributions of different size and then measure the performance, would that be sufficient?

[ezyang]

If I sample from Dirichlet distributions of different size and then measure the performance, would that be sufficient?

SGTM

[nikitaved]

OK, tests are fixed. TODO benchmarks..

[nikitaved]

Hopefully fixed now!

[facebook-github-bot]

@ezyang is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@ezyang merged this pull request in 35cdb78.

[gchanan]

nit: this isn't really accurate (and I realize this was before your change) -- the targets are given however the users give them. They are interpreted as probabilities by default.