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. 
Edit Report
Pub: 27 Feb 2024 19:08 UTC
Views: 80