EnzymeJaX python error

```
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/pypi_benchmark/site-packages/benchmark/utils.py", line 41, in fit
model.fit(dataset, epochs=1, callbacks=[callback])
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/pypi_keras_hub/site-packages/keras_hub/src/utils/pipeline_model.py", line 177, in fit
return super().fit(
^^^^^^^^^^^^
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/pypi_keras/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/__main__/src/enzyme_ad/jax/primitives.py", line 1831, in wrapped
out_shape = jitres.eval_shape(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/__main__/src/enzyme_ad/jax/primitives.py", line 1831, in wrapped
out_shape = jitres.eval_shape(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.bazel/sandbox/processwrapper-sandbox/13615/execroot/__main__/bazel-out/k8-opt/bin/test/keras_test.runfiles/pypi_jax/site-packages/jax/_src/errors.py", line 521, in __init__
f"Attempted boolean conversion of {tracer._error_repr()}."
^^^^^^^^^^^^^^^^^^
AttributeError: Exception encountered when calling BertBackbone.call().
```

x/ref https://github.com/EnzymeAD/Enzyme-JAX/pull/1188

Start with src/enzyme_ad/jax/primitives.py at line 1831 and inspect the recursive wrapped/eval_shape calls shown in the traceback. Compare the failure with the referenced pull request 1188 and the BertBackbone.call path; the issue provides no reproduction steps or acceptance criteria, so the expected fix is not defined.

Written by the indexing model from the issue text.