-
|
This is the testing code for comparing the reshape function in Numpy and Jax. It reshapes 2 The Jax code runs on a V100 card with 30G memory. And the Numpy code runs on Intel(R) Xeon(R) CPU E5-2630 v4 @ 2.20GHz. I'm curious why this happens and how to solve it? |
Beta Was this translation helpful? Give feedback.
Replies: 3 comments 10 replies
-
|
Possible causes:
from functools import partial
import jax
import jax.numpy as jnp
@partial(jax.jit, donate_argnums=0) # use donate argnums to enable inplace update
def modify(a):
return a.at[0].set(10)
x = jnp.arange(4)
y = x.reshape((2, 2))
modify(x)
print(y) # unmodifiedimport numpy as np
x = np.arange(4)
y = x.reshape((2, 2))
x[0] = 10
print(y) # modified |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for pointing this! I have tried the jax.numpy.resize(), but this function gives the similar results. |
Beta Was this translation helpful? Give feedback.
-
|
Hi! - You don't need to worry about it, just jit your function, then XLA will optimize it for you. XLA usually emits bitcast op which is zero cost. |
Beta Was this translation helpful? Give feedback.
-
|
Hi~ Actually, It is a little pity that my total function cannot be jitted together, because the shape of output depend on the input. I found the reshape part of the function occupies a lot of time. So I want to find a way around this problem. |
Beta Was this translation helpful? Give feedback.
-
|
@c-wanen You might try import functools as ft
import jax
@ft.partial(jax.jit, donate_argnums=0, static_argnums=1)
def reshape(a, ord):
return jax.numpy.reshape(a, ord)
for i in range(40):
print(f'---- step {i + 1} -----')
start = time.time()
v = reshape(v, (d, d, d, d))
v.block_until_ready()
end = time.time()
print('time for jit reshape', end - start)
v = reshape(v, (d * d, d * d))
v.block_until_ready() |
Beta Was this translation helpful? Give feedback.
-
|
Hello - this is entirely expected behavior for microbenchmarks of small, individual operations like Also, when writing microbenchmarks like this with JAX, you need to be careful about effects of asynchronous dispatch, data transfer, and other issues. For more information on how to write good microbenchmarks, see FAQ: Benchmarking JAX Code. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for your helpful suggestions~ When I added the block_until_ready() after reshape, the time for reshape is now always the same as 10^-2 second. So the shorter time in the former step is due to the asynchronous dispatch. |
Beta Was this translation helpful? Give feedback.
-
|
There is no way for JAX to achieve this for a stand-alone, eagerly computed |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for your patient clarification! I won't worry about it anymore~ |
Beta Was this translation helpful? Give feedback.
-
|
@jakevdp Why buffer donation requires same shape instead of same size? from functools import partial
from contextlib import contextmanager
import jax
@contextmanager
def timer(msg):
from time import time
t = time()
yield
print(f'{msg}: {time() - t}')
@partial(jax.jit, donate_argnums=0, static_argnums=1)
def reshape(a, ord):
return jax.numpy.reshape(a, ord)
d = 100
v = jax.numpy.zeros((d * d, d * d))
for i in range(3):
print(f'---- step {i + 1} -----')
with timer('reshape time:'):
v = reshape(v, (d, d, d, d))
v.block_until_ready()
v = reshape(v, (d * d, d * d))
v.block_until_ready()And I got warning:
|
Beta Was this translation helpful? Give feedback.
-
|
If you have a question about buffer donation, I'd suggest opening a new discussion. Asking a new question in the answer to another discussion is not likely to get your question the attention it needs. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks! Could I open a feature request issue? |
Beta Was this translation helpful? Give feedback.
-
|
@YouJiacheng Hi Youjia, Thanks for your examples! I think I can understand it better with it. |
Beta Was this translation helpful? Give feedback.
Hello - this is entirely expected behavior for microbenchmarks of small, individual operations like
reshape, though it will not generally result in slower execution when thereshapeis used in more realistic situations (i.e. a sequence of JIT-compiled operations). For information on why this may be expected, see FAQ: Is JAX Faster Than NumPy?.Also, when writing microbenchmarks like this with JAX, you need to be careful about effects of asynchronous dispatch, data transfer, and other issues. For more information on how to write good microbenchmarks, see FAQ: Benchmarking JAX Code.