Add support for passing array attributes via ffi_call and an example demonstrating the different attribute APIs supported by the FFI

[dfm]

The XLA FFI supports a rich interface for specifying "attributes", but this interface hasn't been widely used from the JAX side. We discovered that passing array attributes wasn't supported by the current front end, so I added support for that and took advantage of this change to add a more complete example of the e2e usage of attributes.

The core JAX change that was needed was to add support for passing np.array objects (which aren't hashable!) as keyword arguments to the ffi_call primitive. To do this, I wrap any unhashable keyword arguments in the _HashableByObjectId helper which creates a hashable object based on the id.

"Why not pass these unhashable objects as positional arguments?" you ask. Well, the problem is that we need access to the concrete value of this array when lowering, so we can't have them turning into tracers.

I don't expect this to introduce any real caching issues, but if anyone has concerns or other suggestions for a better API, I'd love to hear them!

[gspschmid]

I don't expect this to introduce any real caching issues, but if anyone has concerns or other suggestions for a better API, I'd love to hear them!

Just to clarify, does this defeat the hashing at the JAX level and would any ffi_call using arrays or dicts hence be uncacheable at the JAX level (but presumably cached fine at the MLIR level)? If I recall correctly, there's a caching mechanism for jitted functions both within JAX itself (on jaxprs?) and another on the emitted StableHLO.