gensbi.flow_matching.utils#

Submodules#

Functions#

expand_tensor_like(input_array, expand_to)

input_array is a 1d vector of length equal to the batch size of expand_to,

unsqueeze_to_match(source, target[, how])

Unsqueeze the source array to match the dimensionality of the target array.

Package Contents#

gensbi.flow_matching.utils.expand_tensor_like(input_array, expand_to)[source]#

input_array is a 1d vector of length equal to the batch size of expand_to, expand input_array to have the same shape as expand_to along all remaining dimensions.

Parameters:

  • input_array (Array) – (batch_size,).

  • expand_to (Array) – (batch_size, …).

Returns:

(batch_size, …).

Return type:

Array

gensbi.flow_matching.utils.unsqueeze_to_match(source, target, how='suffix')[source]#

Unsqueeze the source array to match the dimensionality of the target array.

Parameters:

  • source (Array) – The source array to be unsqueezed.

  • target (Array) – The target array to match the dimensionality of.

  • how (str, optional) – Whether to unsqueeze the source array at the beginning (“prefix”) or end (“suffix”). Defaults to “suffix”.

Returns:

The unsqueezed source array.

Return type:

Array