Fix some type hints, remove unnecessary parameters, and add some asserts to satisfy the type checker.