jraph
jraph copied to clipboard
pad_with_graphs written with numpy
Hello! I was wondering if there is any particular reason that the pad_with_graphs
function uses the numpy library rather than the jax.numpy library. It looks like every numpy function in there can just be replaced with jax.numpy without any issues, but I could be missing something.