docs: def python_type_to_dtype(python_type: type) -> Optional[tf.DType]:
      '''Maps python types to TensorFlow dtypes.'''
      dtype = None
      if python_type == float:
          dtype = tf.float32
      elif python_type == int:
          dtype = tf.int32
      elif python_type == bool:
          dtype = tf.bool
      return dtype