bijx.AutoJacVF

class bijx.AutoJacVF

Bases: Module

Automatic Jacobian computation for vector fields.

Wraps a vector field to automatically compute the trace of its Jacobian, which is needed for the log-density time derivative in continuous normalizing flows. This eliminates the need to manually implement the divergence computation, but is generally inefficient in higher dimensions.

Parameters:

• vector_field_base – Base vector field (t, x, **kwargs) -> dx/dt.

• event_dim – Number of event dimensions for proper shape handling.

Example

>>> auto_vf = AutoJacVF(lambda t, x: -x, event_dim=1)
>>> velocity, log_density_rate = auto_vf(t, x)

__init__(vector_field_base, event_dim=1)