Understanding Performance and Compatibility Issues with jax.numpy Arrays and PRNGs

[Daniel-Haas-B]

Hey!

I hope it's cool if I ask a pretty basic question here. I'm kinda new to jax and hit a snag that's making me scratch my head. Sorry for the long post ahead, but I've got a ton of details and questions all tangled up, so bear with me, please.

Context
I've been diving into using jax for neural quantum states (just a bit of context in case anyone's doing that too). Basically, I'm working with a neural network, like an MLP, and feeding it inputs that I get from Monte Carlo sampling (which matters because we're dealing with PRNGs here). To optimize my network, I use jax.grad for differentiation. But here's the thing: at some point, I need to get the gradient (and Laplacian) of the MLP with respect to its inputs.- I'm using jit compilation and vmap for some of my functions too.

Not too far in, I realized my sampling was using numpy for random numbers, which meant my network's inputs were np arrays. And my network's parameters were np arrays too. But then I learned that for jit and autograd magic in jax to work, I gotta use jax.numpy arrays instead. So, I switched my inputs and parameters to jax.numpy arrays.

Problem (?)
I started by switching the parameters to jnp's, and surprisingly, it didn't really affect my code's speed or results...

Next, I had to take care of my inputs. As jax's PRNGs are slower than numpy's fancy PCG (as far as I know), I tried simply to convert my numpy random array to a jax.numpy array using "jnp.array(rng.normal(loc=state.positions, scale=self.scale))". My code follows all the immutable array rules and doesn't throw errors, but this move slowed my code down by like 4x. So, I tried using jax.random number generator properly, but it's still way slower than the numpy stuff - again, around 4x.

This leaves me with a couple of mysteries:

• How come I didn't hit any errors when I was using numpy arrays with jit and grad? Is jax sneaky and converting things on the fly? Or is it just silent, leaving me to potentially mess up my results without any heads-up?
• Why bother switching to jnp arrays if it's just gonna slow things down, especially in this case? I'm still running things on a CPU, so maybe there's hope for speed gains on a GPU? I did try running my old numpy-based code on a GPU, thinking it might speed things up, but it actually ran slower. And yeah, I double-checked to make sure it was indeed running on the GPU.

Thanks a ton for sticking with my ramble!

[jakevdp]

How come I didn't hit any errors when I was using numpy arrays with jit and grad?

Numpy arrays are silently converted to JAX arrays (and their contents placed on the default device) when passed to a JAX function.

Why bother switching to jnp arrays if it's just gonna slow things down, especially in this case?

It's hard to say without knowing more about your use-case. Can you show a minimal example of the code you're running?

One potential gotcha here is that under JAX's JIT, Python operations (including numpy ops like random.normal) are executed once at trace-time, and their outputs embedded as compile-time constants into the JAX computation. So if you use np.random within JIT, it may not be doing what you expect. As a quick example:

import jax
import numpy as np

np.random.seed(0)

@jax.jit
def f():
  return np.random.normal()

print(f()) # 1.7640524
print(f()) # 1.7640524
print(f()) # 1.7640524
print(f()) # 1.7640524
print(f()) # 1.7640524

Using a compile-time constant will be much faster than generating a new random number at runtime, so this may be why you're seeing such a speedup (but that's only a guess without seeing your code).

[Daniel-Haas-B]

Thank you for the answer! This is a good point, however, I am not using jit for the sampling process - only for the training of the net.
The code is a little dependent on several interconnected scripts, but a simplified depiction of it has the following structure:

I generate several proposals of positions randomly, which are fed as input to a neural net. This process boils down to the following step function but step itself is not jitted. Note that this state_batch object below is a batch of dictionary-like structures. Each element of the batch itself is mutable, but the attributes, like "positions" can be jax.numpy arrays.

def step(self, net, state_batch, seed, batch_size):
    for i in range(batch_size):
        state = state_batch[i - 1]
        next_gen = advance_PRNG_state(seed, state.delta)
        rng = self.rng(next_gen)
        
        # proposals_pos = jnp.array(rng.normal(loc=state.positions, scale=self.scale))
        proposals_pos = rng.normal(loc=state.positions, scale=self.scale)
        log_unif = np.log(rng.random())
        logp_proposal = net.logprob(proposals_pos) # this logprob call makes the forward pass

        # acceptance criterion
        accept = log_unif < logp_proposal - state.logp
        new_positions = proposals_pos if accept else state.positions

        # Create new state
        new_logp = logp_proposal if accept else state.logp
        new_n_accepted = state.n_accepted + accept
        new_delta = state.delta + 1

        state_batch[i] = State(new_positions, new_logp, new_n_accepted, new_delta)

    return state_batch

I only use vmap and jit the forward pass and gradients (or hessian) of the network (with respect to inputs but also, parameters).

In the code below, params is an object of a Parameter class which is again dict-like. I do register the class via
register_pytree_node(Parameter, Parameter.tree_flatten, Parameter.tree_unflatten) to be able to use jax.grad properly.

@jax.jit
def forward(self, x, params):
    """Forward pass"""
    weights = [params.get(f"W{i}") for i in range(len(self._layer_sizes) - 1)]
    biases = [params.get(f"b{i}") for i in range(len(self._layer_sizes) - 1)]

    for i, (w, b) in enumerate(zip(weights, biases)):
        x = self.backend.matmul(x, w) + b
        x = self.activation(self._activations[i])(x)

    return x.squeeze(-1)

@jax.jit
def laplacian(self, x, params):
    hessian_wf = vmap(jax.hessian(self.forward), in_axes=(0, None))
    trace_hessian = vmap(self.backend.trace)
    return trace_hessian(hessian_wf(x, params))

@jax.jit
def grad_param(self, r, params):
    grad_fn = vmap(jax.grad(self.forward, argnums=1), in_axes=(0, None))
    return grad_fn(r, params)  # still a parameter type

@jax.jit
def grad_input(self, r, params):
    grad_wf = jax.grad(self.forward, argnums=0)
    return vmap(grad_wf, in_axes=(0, None))(r, params)

This might not be super useful as you would likely need more details, but thanks for the help anyways.

[jakevdp]

It sounds like you are running code eagerly on CPU – in this case it's not surprising that NumPy would be faster than JAX (see FAQ: Is JAX faster than NumPy?). As for your question of why you would use a JAX PRNG vs a NumPy PRNG, it's mainly due to the fact that the jax PRNG is compatible with JAX transformations and designed for distributed computation (See JAX PRNG design for some background). Again, if you're just running eager computations on CPU, there's no need for this and (as you found) NumPy's PRNG will be faster.

Does that answer your question?

[Daniel-Haas-B]

The PRNG part makes sense. For the sampling process, I am running eager computations. Now, as the jitted functions indicate, I am calling these precompiled functions inside my step function - again, the jitted function "forward" is called by functions inside step.

So as far as I understand, using NumPy's RNG and converting samples to jnp arrays will probably be redundant at this point, or might even hurt if I do it many times and explicitly in my code as I was.

I just did not expect the code to be ~4x slower on a GPU even when letting the NumPy arrays be silently converted. This is probably because the bottleneck of my code is not forward pass/stuff GPU's shine on. Another option is of course I might have implemented some other thing incorrectly, but I guess this is material for another thread...

Thanks for the help.

[jakevdp]

What part is 4x slower? Are you measuring compilation, or device transfer, or runtime? (or maybe all three?)

We have some general guidance on this in FAQ: Benchmarking JAX code.