Can I use JIT from JAX with NetworkX algorithms? For instance, if were to compute the average clustering coefficient for a NetworkX graph object, is it possible to use the @jit decorator to speed up my analysis pipeline?
1 Answers
No, JAX's JIT and other transforms only work with functions implemented via JAX primitives (generally operations defined in jax.lax, jax.numpy, and related submodules). They cannot be used to compile/transform arbitrary Python code.
Related