Allow raise_error to be bool type and make return type more precise.

[copybara-service[bot]]

Allow raise_error to be bool type and make return type more precise.

Prior to this CL, pytype raises an error on this:

from edward2 import maps
def my_custom_function() -> list[int]:
  return maps.robust_map(lambda x: x + 1, range(10))

Or this:

from edward2 import maps
def my_custom_function(raise_error: bool) -> Sequence[int]:
  return maps.robust_map(lambda x: x + 1, range(10), raise_error=raise_error)