<jupyter_start><jupyter_text>It seems that the bottleneck is the function of `forward_rate_constants_func()`<jupyter_code>%timeit sol.equilibrium_constants_func()
%timeit sol.reverse_rate_constants_func()
%timeit sol.wdot_func()
# Test if the AD works properly

%time

TPY_grad = torch.autograd.grad(outputs=sol.wdot.sum(),
                               inputs=TPY,
                               retain_graph=True,
                               create_graph=True,
                               allow_unused=True)[0]

print(TPY_grad.shape)<jupyter_output>CPU times: user 5 µs, sys: 1e+03 ns, total: 6 µs
Wall time: 11.2 µs
torch.Size([3233, 55])
{'fraction_non_alphanumeric': 0.09722222222222222, 'fraction_numerical': 0.023148148148148147, 'mean_word_length': 2.527173913043478, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '308387', 'license': 'MIT', 'n_tokens_mistral': 226, 'n_tokens_neox': 216, 'n_words': 45, 'path': 'test/Profiler_Line_by_Line.ipynb', 'repo_name': 'erwanp/reactorch', 'size': '648'}