Improve device-resident on-device API semantics - #5

Summary

• update matmul on-device behavior to support device-resident JAX inputs cleanly, including fused FP64 return on-device
• enforce consistent input family requirements (both numpy or both JAX) and reject mixed input types with a clear error
• normalize on-device output type to match input family for predictable API behavior
• extend CPU validation checks to cover JAX-input on-device paths and mixed-type rejection
• add benchmarks/tpu_devres_bench.py for TPU device-resident API benchmarking

Validation notes

CPU workflow coverage now includes:

• JAX-input fused/ondevice/host accumulation paths
• output type assertions for JAX inputs
• mixed input-type rejection behavior

Notes

• fused on-device mode still requires jax_enable_x64=True
• this PR focuses on API semantics and transfer behavior; TPU perf validation remains benchmark-driven

Made with Cursor