* Fixed incorrectly formatted docstrings.
* Fixed use of generic asserts.
* Changed to use absltest to be able to use more specific assert variants.

PiperOrigin-RevId: 224189879