This module is largely modeled after https://github.com/facebookresearch/flow_matching, original copyright (c) Meta Platforms, Inc. and affiliates.