enablxe_x64() fails with jacrev, bug or am I using it wrong?

[Binbose]

Hello, I have a large codebase that runs in float32, but I want to run one operation in float64. It looks like jax.experimental.enable_x64() does exactly what I want, however it fails as soon as I try to backward mode autodiff it (but it works in forward mode). I am not sure if I am using it wrong or if it is a bug. Here is a minimal reproducing example:

def f(x):
    with jax.experimental.enable_x64():
        B = jnp.linalg.det(x.astype(jnp.float64)) # <- fails also for inv, norm etc. but sometimes with different error messages
    return B 

x = jax.random.uniform(key=jax.random.PRNGKey(1), shape=(4,4))


jax.jacfwd(f)(x) # <- works

# Output:
# Array([[ 0.09959597, -0.04547443,  0.01634833, -0.07254504],
#       [ 0.04057878, -0.00533076,  0.00425937, -0.0311282 ],
#       [-0.25385   ,  0.11323178, -0.02177901,  0.1850684 ],
#       [-0.04267051,  0.01776768, -0.0098353 ,  0.03706509]],      dtype=float32)

jax.jacrev(f)(x) # <- fails

---------------------------------------------------------------------------
MLIRError                                 Traceback (most recent call last)
File ~/anaconda3/envs/qmace/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:821, in lower_jaxpr_to_module(module_name, jaxpr, ordered_effects, backend_or_name, platforms, axis_context, name_stack, donated_args, replicated_args, arg_shardings, result_shardings, arg_names, result_names, num_replicas, num_partitions, all_default_mem_kind, lowering_parameters)
    [820](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:820) try:
--> [821](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:821)   if not ctx.module.operation.verify():
    [822](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:822)     module_string = module_to_string(ctx.module)

MLIRError: Verification failed:
error: "jit(triangular_solve)/jit(main)/triangular_solve[left_side=True lower=False transpose_a=True conjugate_a=False unit_diagonal=False]"("/tmp/ipykernel_13577/3195055520.py":4:0): 'stablehlo.triangular_solve' op requires the same element type for all operands and results
 note: "jit(triangular_solve)/jit(main)/triangular_solve[left_side=True lower=False transpose_a=True conjugate_a=False unit_diagonal=False]"("/tmp/ipykernel_13577/3195055520.py":4:0): see current operation: %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {left_side = true, lower = false, transpose_a = #stablehlo<transpose TRANSPOSE>, unit_diagonal = false} : (tensor<4x4xf64>, tensor<4x4xf32>) -> tensor<4x4xf32>

The above exception was the direct cause of the following exception:

JaxStackTraceBeforeTransformation         Traceback (most recent call last)
File ~/anaconda3/envs/qmace/lib/python3.10/runpy.py:196, in _run_module_as_main()
    [195](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:195)     sys.argv[0] = mod_spec.origin
--> [196](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:196) return _run_code(code, main_globals, None,
    [197](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:197)                  "__main__", mod_spec)

File ~/anaconda3/envs/qmace/lib/python3.10/runpy.py:86, in _run_code()
     [79](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:79) run_globals.update(__name__ = mod_name,
     [80](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:80)                    __file__ = fname,
     [81](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:81)                    __cached__ = cached,
   (...)
     [84](https://file+.vscode-resource.vscode-cdn.net/home/luca/Projects/vqmc/qmace/~/anaconda3/envs/qmace/lib/python3.10/runpy.py:84)                    __package__ = pkg_name,
...
  }) : () -> () loc(#loc)
}) {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} : () -> () loc(#loc)
#loc1 = loc("/tmp/ipykernel_13577/3195055520.py":4:0)
#loc2 = loc("jit(triangular_solve)/jit(main)/triangular_solve[left_side=True lower=False transpose_a=True conjugate_a=False unit_diagonal=False]"(#loc1))

[jakevdp]

jax.experimental.enable_x64 will not work in all contexts, and you've hit one context where it is known not to work. There is no workaround; the mechanism it uses is fundamentally flawed, which is one reason it's remained in jax.experimental.