Neil Kichler

Currently, the basic example from the README.md does not run for versions of JAX>=4.25.0. This is due to the removal of the count attribute in the jax.core.Var class and its...

In the README example, Pyright complains about the arguments to taylor_bounds since the types of the Callable are restricted to jnp.ndarray and not types.NDArraylike which also supports int, float etc.

Fix register index bug in mma.sync.aligned.m16n8k16

2

In `include/cutlass/arch/mma_sm90.h` the ptx instruction `mma.sync.aligned.m16n8k16` has a typo in the variable for `%5` which should be `A[1]` but is currently `A[2]` (and hence using `A[2]` twice and `A[1]` not...