EnzymeJaX python error

[wsmoses]

  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 #1188