jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 119, in <module>
main(key)
File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 94, in main
trainer.train()
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 262, in train
loss, model, opt_state = make_step(model, seq, label, pad_mask, rndm_n, rndm_k,
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 206, in __call__
return self._call(False, args, kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_module.py", line 875, in __call__
return self.__func__(self.__self__, *args, **kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 200, in _call
out = self._cached(dynamic_donate, dynamic_nodonate, static)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 90, in make_step
loss, grads = compute_loss(model, x, y, pad_mask, n, k, num_classes, keys)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 66, in compute_loss
pred_y = jax.vmap(forward, in_axes=(None, 0, 0, 0, 0, 0))(model, x, pad_mask, n, k, keys) # (batch_size, seqlen, num_classes)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 49, in iters_fwd
output, _ = model(input_arr, k, pad_mask=pad_mask, key=key)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/model/react.py", line 141, in __call__
output = self.iterate_for_steps(interim_thought, pad_mask, iters_to_do, input_arr, key) # (batch, seqlen, bottleneck)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/model/react.py", line 121, in iterate_for_steps
iters_to_do = iters_to_do.astype(int).item()
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 733, in meth
return getattr(self.aval, name).fun(self, *args, **kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 81, in _item
return int(a)
jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.
The error occurred while tracing the function __call__ at /Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py:37 for jit.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 119, in <module>
main(key)
File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 94, in main
trainer.train()
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 262, in train
loss, model, opt_state = make_step(model, seq, label, pad_mask, rndm_n, rndm_k,
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 206, in __call__
return self._call(False, args, kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_module.py", line 875, in __call__
return self.__func__(self.__self__, *args, **kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 200, in _call
out = self._cached(dynamic_donate, dynamic_nodonate, static)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 90, in make_step
loss, grads = compute_loss(model, x, y, pad_mask, n, k, num_classes, keys)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 66, in compute_loss
pred_y = jax.vmap(forward, in_axes=(None, 0, 0, 0, 0, 0))(model, x, pad_mask, n, k, keys) # (batch_size, seqlen, num_classes)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 49, in iters_fwd
output, _ = model(input_arr, k, pad_mask=pad_mask, key=key)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/model/react.py", line 141, in __call__
output = self.iterate_for_steps(interim_thought, pad_mask, iters_to_do, input_arr, key) # (batch, seqlen, bottleneck)
File "/Users/neel/Documents/research/ReAct_Jax/ReAct/model/react.py", line 121, in iterate_for_steps
iters_to_do = iters_to_do.astype(int).item()
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 733, in meth
return getattr(self.aval, name).fun(self, *args, **kwargs)
File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 81, in _item
return int(a)
jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.
The error occurred while tracing the function __call__ at /Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py:37 for jit.