jax icon indicating copy to clipboard operation
jax copied to clipboard

Published 20 hours ago • verified publisher icon google

simplify slicing jaxprs a little

Open mattjj opened this issue 2 years ago • 0 comments

Before: image

After: image

No more lt, add, select, or convert_element_type!

This might actually improve compilation time a bit; on CPU, the tiny benchmark added went from ~15.5ms to 13ms on my machine.

mattjj avatar Aug 12 '22 02:08 mattjj