Metadata-Version: 2.5
Name: cutejax
Version: 0.1.0
Summary: Call CuTe DSL kernels from JAX via TVM-FFI target registration
Author-email: Emily Shepperd <nshepperd@gmail.com>
License-Expression: MIT
License-File: LICENSE
Requires-Python: >=3.11
Requires-Dist: jax
Requires-Dist: jax-tvm-ffi
Requires-Dist: nvidia-cutlass-dsl
Provides-Extra: cu13
Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.7.0; extra == 'cu13'
Description-Content-Type: text/markdown

# cutejax
 cutedsl jax wrapper that adds some useful features.
 Mostly written by Fable, so *caveat claudor*. But it seems to work.

## 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.Module`s) 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](examples/fa4_attention.py).
- **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.
- **A compile watchdog.** Turns a runaway trace into an error instead of letting it
  eat all your RAM.
- **Convenience.** `cutejax.Module` registers with both jax's and cute's pytree
  registries; `cutejax.cond` gives you data-dependent branches in undecorated code.
- **Works on older jax.** Tested on 0.8.2 through 0.11.0.

## When to use the built-in instead

- You only pass arrays in and get arrays out — then cutejax buys you nothing.
- You need non-row-major layouts. `cutlass_call`'s `TensorSpec(mode=...)` can remap
  dimensions; cutejax can't, because the XLA FFI boundary doesn't carry strides,
  and rejects those layouts up front rather than reading them wrong.
- You'd rather depend only on things NVIDIA ships and supports. cutejax needs
  `jax-tvm-ffi`, and is one person's side project.
