[AudioLM] Graph break: const method call float.is_integer

[ezyang]

🐛 Describe the bug

Discovered while compiling https://github.com/lucidrains/audiolm-pytorch/

[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] Graph break: const method call float.is_integer from user code at:
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/data/users/ezyang/audiolm-pytorch/audiolm_pytorch/soundstream.py", line 830, in forward
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     x = self.encoder_attn(x)
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     return forward_call(*args, **kwargs)
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/data/users/ezyang/audiolm-pytorch/audiolm_pytorch/soundstream.py", line 436, in forward
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     x = attn(x, attn_bias = attn_bias) + x
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     return forward_call(*args, **kwargs)
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/local_attention-1.9.0-py3.10.egg/local_attention/transformer.py", line 106, in forward
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     out = self.attn_fn(q, k, v, mask = mask, attn_bias = attn_bias)
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     return forward_call(*args, **kwargs)
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/local_attention-1.9.0-py3.10.egg/local_attention/local_attention.py", line 126, in forward
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     (needed_pad, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v))
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/local_attention-1.9.0-py3.10.egg/local_attention/local_attention.py", line 126, in <lambda>
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     (needed_pad, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v))
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/local_attention-1.9.0-py3.10.egg/local_attention/local_attention.py", line 37, in pad_to_multiple
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     if m.is_integer():
[2024-03-06 11:43:47,766] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] 

Should be pretty simple to fix.

Full repro code: https://gist.github.com/ezyang/64c24c9fc5529f3afed4ee4266f6adc5

Versions

main

cc @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng