AttributeError: DynamicJaxprTracer has no attribute block_until_ready? How to Write xla.apply_primitive... #17190
Replies: 1 comment 8 replies
-
|
It's hard to know what's going on without seeing the code, but I'd suggest removing |
Beta Was this translation helpful? Give feedback.
-
from functools import partial
import numpy as np
from jax import core, numpy as jnp
from jax.interpreters import xla, mlir
from jax._src import abstract_arrays
import jaxlib.mlir.ir as ir
from jaxlib.hlo_helpers import custom_call
from .register_custom_calls import default_layouts, gpu_ops
import jax
all = ["LIF_cuda"]
_LIF_fwd_p = core.Primitive("LIF_fwd")
_LIF_fwd_p.multiple_results = True
_LIF_fwd_p.def_impl(partial(xla.apply_primitive, _LIF_fwd_p))
def LIF_fwd(inp, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum):
(spike, membrane) = _LIF_fwd_p.bind(inp, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum)
return (spike, membrane), (spike, membrane, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem)
_LIF_bwd_p = core.Primitive("LIF_bwd")
_LIF_bwd_p.multiple_results = True
_LIF_bwd_p.def_impl(partial(xla.apply_primitive, _LIF_bwd_p))
def LIF_bwd(res, output_grad):
spike, membrane, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem = res
spike_dot, membrane_dot = output_grad[0], output_grad[1]
from LIF import global_surr
surrogates = global_surr(membrane, theta)
not_clipped = jnp.ones_like(surrogates)
post_grad = surrogates * spike_dot + membrane_dot
alpha_memSub = alpha * mem_sub
post_mem = membrane - mem_sub.reshape(-1, 1) * spike
grad_input, grad_alpha = _LIF_bwd_p.bind(surrogates, post_grad, not_clipped, alpha, alpha_memSub, post_mem,
mem_init, mem_sub)
import jax
import jaxlib
print("jax version: ", jax.__version__, "\t", "jaxlib version: ", jaxlib.__version__)
print("grad_input: ", grad_input.sum(), "\t", "grad_alpha: ", grad_alpha.sum())
grad_init = alpha * grad_input[:, 0]
return (grad_input, grad_init, grad_alpha, None, None, None, None, None)The same code only relies on different versions of jaxlib and jaxlib: from functools import partial
import numpy as np
from jax import core, numpy as jnp
from jax.interpreters import xla, mlir
from jax._src import abstract_arrays
import jaxlib.mlir.ir as ir
from jaxlib.hlo_helpers import custom_call
from .register_custom_calls import default_layouts, gpu_ops
import jax
all = ["LIF_cuda"]
_LIF_fwd_p = core.Primitive("LIF_fwd")
_LIF_fwd_p.multiple_results = True
_LIF_fwd_p.def_impl(partial(xla.apply_primitive, _LIF_fwd_p))
def LIF_fwd(inp, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum):
(spike, membrane) = _LIF_fwd_p.bind(inp, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum)
return (spike, membrane), (spike, membrane, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem)
_LIF_bwd_p = core.Primitive("LIF_bwd")
_LIF_bwd_p.multiple_results = True
_LIF_bwd_p.def_impl(partial(xla.apply_primitive, _LIF_bwd_p))
def LIF_bwd(res, output_grad):
spike, membrane, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem = res
spike_dot, membrane_dot = output_grad[0], output_grad[1]
from LIF import global_surr
surrogates = global_surr(membrane, theta)
not_clipped = jnp.ones_like(surrogates)
post_grad = surrogates * spike_dot + membrane_dot
alpha_memSub = alpha * mem_sub
post_mem = membrane - mem_sub.reshape(-1, 1) * spike
post_mem.block_until_ready()
grad_input, grad_alpha = _LIF_bwd_p.bind(surrogates, post_grad, not_clipped, alpha, alpha_memSub, post_mem,
mem_init, mem_sub)
import jax
import jaxlib
print("jax version: ", jax.__version__, "\t", "jaxlib version: ", jaxlib.__version__)
print("grad_input: ", grad_input.sum(), "\t", "grad_alpha: ", grad_alpha.sum())
grad_init = alpha * grad_input[:, 0]
return (grad_input, grad_init, grad_alpha, None, None, None, None, None)` |
Beta Was this translation helpful? Give feedback.
-
|
(I edited your code - in the future you can use triple ticks to mark codeblocks – hopefully I got the indentations correct, I was a bit thrown off by the top-level imports within function bodies) |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for pasting the code. I don't entirely understand what's going on: I don't see any lowering rule registered for If using |
Beta Was this translation helpful? Give feedback.
-
|
If you want to make progress here, I'd suggest opening a bug with a minimal reproducible example: in other words, a single block of code that someone could paste into an editor and execute to see the outputs you're seeing. |
Beta Was this translation helpful? Give feedback.
-
|
Thank you for adjusting the code format, which makes it look clearer. I will try to reproduce this bug with minimal reproducible example and paste the code here. |
Beta Was this translation helpful? Give feedback.
Uh oh!
There was an error while loading. Please reload this page.
-
What if I need to define a calculated backpropagation and perform a series of calculations before using CUDA in the backward function? The calculation is correct in non jit cases because I used block_until_readyy.
Beta Was this translation helpful? Give feedback.
All reactions