Decouple from JAX

[joelberkeley]

Hi, is there any interest in pulling out the non-JAX specific stuff, such as autodiff of (Stable)HLO? It would reduce cognitive overhead for users, and simplify builds.

I frequently have trouble debugging build issues regarding JAX, and you're overriding the XLA commit used by JAX, so I'm not even sure I'm using anything from JAX.

[mofeing]

So currently we use jax(lib) for the following:

• FFI helpers and CUDA BLAS handle pool

  Enzyme-JAX/src/enzyme_ad/jax/xla_ffi/cuda/blas.cc

  Lines 9 to 10 in bc8343f

  	#include "jaxlib/ffi_helpers.h"
  	#include "jaxlib/gpu/blas_handle_pool.h"

• Triton protobuf

  Enzyme-JAX/src/enzyme_ad/jax/Passes/RaiseTritonCustomCallPass.cpp

  Line 30 in bc8343f

  #include "jaxlib/gpu/triton.pb.h"

• Some CUDA custom call implementations (specifically CuSolver, but maybe more?)

  Enzyme-JAX/src/enzyme_ad/jax/Passes/LowerEnzymeXLALapack.cpp

  Line 223 in bc8343f

  // emit `stablehlo.custom_call` to `@cusolver_geqrf_ffi` kernel from jaxlib

Removing depende on CUDA custom calls is doable but we're on some deadlines so it'll need to wait for a couple of weeks.

@avik-pal what do you think about the other points?