> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.flow_matching.time_shift_utils

## Module Contents

### Functions

| Name                                                                                                                                       | Description                                                                 |
| ------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------- |
| [`compute_density_for_timestep_sampling`](#nemo_automodel-components-flow_matching-time_shift_utils-compute_density_for_timestep_sampling) | Sample timesteps from different distributions for better training coverage. |
| [`get_flow_match_loss_weight`](#nemo_automodel-components-flow_matching-time_shift_utils-get_flow_match_loss_weight)                       | Compute loss weights for flow matching based on sigma values.               |
| [`time_shift`](#nemo_automodel-components-flow_matching-time_shift_utils-time_shift)                                                       | Convert timesteps to sigmas with sequence-length-aware shifting.            |

### API

```python
nemo_automodel.components.flow_matching.time_shift_utils.compute_density_for_timestep_sampling(
    weighting_scheme: str,
    batch_size: int,
    logit_mean: float = 0.0,
    logit_std: float = 1.0,
    mode_scale: float = 1.29
)
```

Sample timesteps from different distributions for better training coverage.

**Parameters:**

"uniform", "logit\_normal", or "mode"

number of samples to generate

mean for logit-normal distribution

std for logit-normal distribution

scale for mode-based sampling

**Returns:**

Tensor of shape (batch\_size,) with values in \[0, 1]

```python
nemo_automodel.components.flow_matching.time_shift_utils.get_flow_match_loss_weight(
    sigma: torch.Tensor,
    shift: float = 3.0
)
```

Compute loss weights for flow matching based on sigma values.

Higher sigma (more noise) typically gets higher weight.

**Parameters:**

sigma values in range \[0, 1]

weight scaling factor

**Returns:**

Loss weights with same shape as sigma

```python
nemo_automodel.components.flow_matching.time_shift_utils.time_shift(
    t: torch.Tensor,
    image_seq_len: int,
    shift_type: str = 'constant',
    base_shift: float = 0.5,
    max_shift: float = 1.15,
    constant: float = 3.0
)
```

Convert timesteps to sigmas with sequence-length-aware shifting.

**Parameters:**

timesteps in range \[0, 1]

number of tokens (frames \* height \* width / patch\_size^2)

"linear", "sqrt", or "constant"

base shift for linear mode

max shift for linear mode

shift value for constant mode (default 3.0 matches Pika)

**Returns:**

sigma values for noise scheduling