jax

48c766cd - Update error message in a few places where we say `operand shape` when we mean `operand type`.