Miscompiles for VJP when using jax.lax.associative_scan

I was playing around with `enzyme_jax_ir` and found some compile errors which come up when using `jax.lax.associative_scan` where either the whole scan or the scan body is annotated with `enzyme_jax_ir`:

```python
import jax
import jax.numpy as jnp
from enzyme_ad.jax import enzyme_jax_ir

def linear_rnn_cell_companion_op(x, y):
a1, b1 = x
a2, b2 = y
return a1 * a2, a2 * b1 + b2

def linear_rnn(a, x):
_, out = jax.lax.associative_scan(linear_rnn_cell_companion_op, (a, x), axis=-1)
return out

@enzyme_jax_ir()
def fused_linear_rnn_cell_companion_op(x, y):
a1, b1 = x
a2, b2 = y
return a1 * a2, a2 * b1 + b2

@enzyme_jax_ir()
def fused_linear_rnn(a, x):
_, out = jax.lax.associative_scan(linear_rnn_cell_companion_op, (a, x), axis=-1)
return out

def linear_rnn_fused_op(a, x):
_, out = jax.lax.associative_scan(
fused_linear_rnn_cell_companion_op, (a, x), axis=-1
)
return out

def make_vjp_fn(fn):
@jax.jit
def _vjp_fn(g, *args):
_, vjp = jax.vjp(fn, *args)
return vjp(g)

return _vjp_fn

key = jax.random.PRNGKey(0)
B, C = 8, 4

T = 512
a = jax.random.uniform(key, (B, C, T), minval=0, maxval=1.0)
x = jax.random.normal(key, (B, C, T))
g = jnp.ones_like(x)

# works
y = linear_rnn(a, x)
g = jnp.ones_like(y)
vjp_fn = make_vjp_fn(linear_rnn)
print(vjp_fn.lower(g, a, x).as_text())

# works
y = fused_linear_rnn_cell_companion_op((a, x), (a, x))[0]
g = jnp.ones_like(y)
vjp_fn = make_vjp_fn(lambda x, y: fused_linear_rnn_cell_companion_op(x, y)[0])
print(vjp_fn.lower(g, (a, x), (a, x)).as_text())

# works
y = fused_linear_rnn_cell_companion_op((a, x), (a, x))[1]
g = jnp.ones_like(y)
vjp_fn = make_vjp_fn(lambda x, y: fused_linear_rnn_cell_companion_op(x, y)[1])
print(vjp_fn.lower(g, (a, x), (a, x)).as_text())

# TODO: breaks - 'arith.addf' op requires the same type for all operands and results
y = fused_linear_rnn(a, x)
g = jnp.ones_like(y)
vjp_fn = make_vjp_fn(fused_linear_rnn)
print(vjp_fn.lower(g, a, x).as_text())

# TODO: breaks - Dialect `tensor' not found for custom op 'tensor.empty'
y = linear_rnn_fused_op(a, x)
g = jnp.ones_like(y)
vjp_fn = make_vjp_fn(linear_rnn_fused_op)
print(vjp_fn.lower(g, a, x).as_text())

```

```
loc("x"): error: 'arith.addf' op requires the same type for all operands and results
Traceback (most recent call last):
File "/home/falkaer/projects/latent-features/jax_poc.py", line 80, in
print(vjp_fn.lower(g, a, x).as_text())
File "/home/falkaer/projects/latent-features/jax_poc.py", line 40, in _vjp_fn
_, vjp = jax.vjp(fn, *args)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Failed to parse pipeline
Pipeline failed:
'arith.addf' op requires the same type for all operands and results
```

```
Traceback (most recent call last):
File "/home/falkaer/projects/latent-features/jax_poc.py", line 84, in
y = linear_rnn_fused_op(a, x)
File "/home/falkaer/projects/latent-features/jax_poc.py", line 31, in linear_rnn_fused_op
_, out = jax.lax.associative_scan(
File "/home/falkaer/projects/latent-features/.devenv/state/venv/lib/python3.11/site-packages/enzyme_ad/jax/primitives.py", line 1818, in wrapped
out_flat = ffi_call(
File "/home/falkaer/projects/latent-features/.devenv/state/venv/lib/python3.11/site-packages/enzyme_ad/jax/primitives.py", line 1296, in ffi_call
return _enzyme_primal_p.bind(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: jaxlib.mlir._mlir_libs._site_initialize..MLIRError: Unable to parse module assembly:
error: "-":7:10: Dialect `tensor' not found for custom op 'tensor.empty'
note: "-":7:10: Registered dialects: arith, builtin, cf, chlo, func, gpu, llvm, math, memref, mhlo, nvgpu, nvvm, scf, sdy, stablehlo, vector ; for more info on dialect registration see https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management
```

Run the provided reproducer in jax_poc.py, focusing on the two failing VJP cases using jax.lax.associative_scan and enzyme_jax_ir. Inspect enzyme_ad/jax/primitives.py around ffi_call and _enzyme_primal_p.bind, then verify that both examples lower without the arith.addf type error or missing tensor dialect error.

Written by the indexing model from the issue text.