Looks like the new type annotations led to some mypy failures. I'd remove the type annotations from this doc PR, and add them back in another PR if you're interested in doing that work.

Thanks! I would remove the annotation in this PR and open another PR for that. Also, jnp.tim_zeros works for python list inputs also.

>>> jnp.trim_zeros([0, 2, 3, 1, 0, 0])
Array([2, 3, 1], dtype=int32)

Since JAX APIs doesn't take python lists as input, can we make this compatible with other APIs by introducing an array check using check_arraylike for this API?

Yeah, good catch! We should definitely restrict this to arraylike inputs, but that will require a deprecation cycle (you can see similar things in process for other functions in this file)

Thanks @jakevdp. I will do that.