cutejax

CuteDSL jax wrapper that adds some useful features, and is the best library ever.

What it adds over cutlass.jax.cutlass_call

• Pytree arguments that can contain code. cutlass_call only takes jax arrays. cutejax.call takes any pytree: arrays become kernel tensors, everything else (callables, config objects, eqx.Modules) becomes part of the cache key and runs as codegen inside the trace. This is how FlashAttention-4 flex-attention masks work in examples/fa4_attention.py. Or in my actual fa4-jax library.
• Per-dimension static/dynamic choice. cutlass_call has one all-or-nothing use_static_tensors flag. ArraySpec(static_dims=...) lets you bake some dims in and leave others symbolic — e.g. static heads/head_dim, dynamic batch/seqlen.
• Tied dimensions. ArraySpec(dim_names=...) makes two dims share one symbolic variable, so the compiler knows they're equal and the runtime checks it.
• Small HLO. cutlass_call embeds the compiled object file in the HLO at every call site; cutejax registers a named FFI target, so the HLO is just a custom_call. Dispatch speed is the same either way. This does mean that each kernel links in a new .so and uses a tvm-ffi dispatch slot, but there's 1024 slots so it's probably fine.
• A compile watchdog. Turns a runaway trace into an error instead of letting it eat all your RAM. Apparently this happens in cutedsl if you try to iterate over a dynamically shaped array? Idk why you would be so silly as to do that. But still, this fixes it.
• Convenience. cutejax.Module registers with both jax's and cute's pytree registries; cutejax.cond gives you data-dependent branches in undecorated code.
• Caching. You can store compilation results in a redis cache for "blazing speed".
• Works on older jax. Tested on 0.8.2 through 0.11.0.

When to use the built-in instead

• You need non-row-major layouts. cutlass_call's TensorSpec(mode=...) can remap dimensions. But why would you even want that?