Selecting arrays by index in batched computations #21010
-
|
I have an application (sparse mixture of experts model) where I need to select a matrix from a batch of matrices to use in batched matrix multiplication. For example: I have a batch of inputs Basically my question is: is there some way to either convert the vmapped implementation so that it does not instantiate the large Example code:import jax
import jax.numpy as jnp
weights = jnp.zeros((8, 128, 512))
xs = jnp.zeros((256, 128))
idx = jnp.arange(len(xs)) % len(weights)
def broadcast(weights, xs, idx):
return jnp.squeeze(xs[:, None] @ weights[idx], -2)
def vmapped(weights, xs, idx):
return jax.vmap(lambda xi, i: xi @ weights[i])(xs, idx)
def scanned(weights, xs, idx):
return jax.lax.map(lambda a: a[0] @ weights[a[1]], (xs, idx))
print("<===== broadcasted =====>")
print(jax.make_jaxpr(broadcast)(weights, xs, idx))
print()
print("<======= vmapped =======>")
print(jax.make_jaxpr(vmapped)(weights, xs, idx))
print()
print("<======= scanned =======>")
print(jax.make_jaxpr(scanned)(weights, xs, idx))
print()Output jaxpr: |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 4 replies
-
|
You may be able to rely on the compiler to fuse operations. Even if the jaxpr indicates an intermediate value of a particular shape, it doesn't necessarily mean that the compiled operation will instantiate that intermediate value. For example, here's the compiled HLO produced by your print(jax.jit(broadcast).lower(weights, xs, idx).compile().as_text())I believe this means the reduction will be fused so that the implied |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for your quick reply. Interesting to see how it differs on CPU vs GPU. Would it be correct to interpret that Would that imply that the sizes of arrays in the HLO code are symbolic and not concrete? |
Beta Was this translation helpful? Give feedback.
-
|
I've just checked with a memory profiler and it seems that xla does not create the |
Beta Was this translation helpful? Give feedback.
-
Yes, precisely. The HLO appearing within a fusion represents the logical operations that will be performed, but you should think of that codeblock as being executed as a single kernel on the device, and intermediate values will not necessarily be materialized (though the compiler can choose to do so if it deems it to be the most efficient approach given available hardware). You can get a sense of what the compiler is doing by using the There's no way that I know of to enforce non-materialization of any particular intermediate value, unfortunately. If you have very specific logic that you need to control at that level, one option is to write the kernel yourself at a lower level using something like pallas. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks! 🙏 |
Beta Was this translation helpful? Give feedback.
You may be able to rely on the compiler to fuse operations. Even if the jaxpr indicates an intermediate value of a particular shape, it doesn't necessarily mean that the compiled operation will instantiate that intermediate value.
For example, here's the compiled HLO produced by your
broadcastfunction on a T4 GPU: