jraph icon indicating copy to clipboard operation
jraph copied to clipboard

pad_with_graphs written with numpy

Open GrantMcConachie opened this issue 7 months ago • 1 comments

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.

GrantMcConachie avatar Jul 05 '24 18:07 GrantMcConachie