Compiling a Bilt model with JAX¶

This example evaluates the same Bilt model with the Python and JAX backends and compares their performance on a large array calculation.

The model represents a simplified cruise-impact calculation. Cruise conditions feed atmosphere, aerodynamic, propulsion, emissions, and illustrative climate-effect calculations before combining into a final output.

The climate-effect relationships are deliberately illustrative rather than a validated aircraft climate model. The point of the example is the shape and size of the numerical calculation.

flowchart TD
    I[Operating conditions] --> A[Atmosphere]
    I --> W[Weight]
    A --> Q[Dynamic pressure]
    W --> L[Lift coefficient]
    Q --> L

    Q --> PD[Parasite drag]
    Q --> ID[Induced drag]
    Q --> CD[Compressibility drag]
    Q --> TD[Trim drag]
    L --> ID
    L --> TD

    PD --> D[Total drag]
    ID --> D
    CD --> D
    TD --> D

    D --> P[Propulsive power]
    A --> P
    P --> F[Fuel flow]

    F --> CO2[CO2]
    F --> NOX[NOx]
    F --> H2O[Water]
    F --> SOOT[Soot]

    CO2 --> C1[CO2 effect]
    NOX --> C2[NOx effect]
    H2O --> C3[Water effect]
    H2O --> C4[Contrail effect]
    SOOT --> C5[Soot effect]
    SOOT --> C4

    C1 --> T[Total climate effect]
    C2 --> T
    C3 --> T
    C4 --> T
    C5 --> T

The model¶

The full model contains a number of small transforms. For example:

from typing import Annotated

from bilt import Quantity, transform
from bilt.units.dimensions import DIMENSIONLESS, FORCE, LENGTH, MASS, TIME


Density = Annotated[Quantity, MASS / LENGTH**3]
Speed = Annotated[Quantity, LENGTH / TIME]
Pressure = Annotated[Quantity, FORCE / LENGTH**2]
Area = Annotated[Quantity, LENGTH**2]
Force = Annotated[Quantity, FORCE]
Ratio = Annotated[Quantity, DIMENSIONLESS]


@transform
def dynamic_pressure(density: Density, true_airspeed: Speed) -> Pressure:
    return 0.5 * density * true_airspeed**2


@transform
def induced_drag(
    dynamic_pressure: Pressure,
    wing_area: Area,
    lift_coefficient: Ratio,
    induced_drag_factor: Ratio,
) -> Force:
    return dynamic_pressure * wing_area * induced_drag_factor * lift_coefficient**2

The system and transforms use scalar quantity declarations. They accept scalar quantities or compatible unit-carrying tensors, so the model does not prescribe the shape of the ensemble. Efficiencies, Mach number, aerodynamic coefficients, and emission indices are dimensionless quantities supplied with UNITLESS.

The atmosphere uses absolute temperatures and relative temperature offsets. Gravity, lapse rate, gas constant, reference pressure, and fuel heating value carry units in the full model. Absolute temperatures are explicitly expressed in kelvin before their numerical values enter the empirical correlations.

Emission rates and the illustrative weighted-emission score have units of kg/s. The score uses arbitrary weights; it is not a validated climate metric or a CO2-equivalent emission rate. The soot-dependent factor uses an explicit inverse mass-rate coefficient so its correction is dimensionless.

The complete model is executed as part of this page, but its remaining definitions are hidden to keep the worked example focused.

# Example model setup is executed here but hidden from the rendered page.

Build a large ensemble¶

We evaluate the model over 262,144 (2**18) cruise conditions.

The inputs use Bilt tensors over a named scenario dimension:

import numpy as np

from bilt import Quantity, dim, tensor
from bilt.units import UNITLESS, kelvin, kilogram, meter, second

scenario = dim("scenario")
scenario_count = 2**18
rng = np.random.default_rng(42)

inputs = {
    "altitude": tensor(
        rng.uniform(8_000.0, 12_000.0, scenario_count).astype(np.float32),
        over=scenario,
        unit=meter,
    ),
    "mach": tensor(
        rng.uniform(0.70, 0.88, scenario_count).astype(np.float32),
        over=scenario,
        unit=UNITLESS,
    ),
    "mass": tensor(
        rng.uniform(50_000.0, 80_000.0, scenario_count).astype(np.float32),
        over=scenario,
        unit=kilogram,
    ),
    "temperature_offset": tensor(
        rng.normal(0.0, 8.0, scenario_count).astype(np.float32),
        over=scenario,
        unit=kelvin,
    ),
    "wing_area": Quantity(float(np.float32(122.6)), meter**2),
    "zero_lift_drag_coefficient": Quantity(float(np.float32(0.021)), UNITLESS),
    "induced_drag_factor": Quantity(float(np.float32(0.045)), UNITLESS),
    "propulsive_efficiency": Quantity(float(np.float32(0.38)), UNITLESS),
    "nox_emission_index": Quantity(float(np.float32(0.015)), UNITLESS),
    "soot_emission_index": Quantity(float(np.float32(0.00003)), UNITLESS),
}

print(f"Scenarios: **{scenario_count:,}**")

Scenarios: 262,144

The scalar aircraft parameters are broadcast across the scenario dimension.

Run with the Python backend¶

Python is Bilt's default backend, so it does not need to be selected explicitly:

python_model = CruiseImpact(**inputs)

with cruise_impact:
    python_result = python_model.total_climate_effect

print(f"Output shape: `{python_result.values.shape}`")
print(f"Output unit: `{python_result.unit}`")

Output shape: (262144,) Output unit: kilogram / second

Run with the JAX backend¶

The model and inputs do not change. Only the backend selection changes:

from bilt import JAX_BACKEND

jax_model = CruiseImpact(**inputs)

with JAX_BACKEND, cruise_impact:
    jax_result = jax_model.total_climate_effect

print(f"Output shape: `{jax_result.values.shape}`")
print(f"Output unit: `{jax_result.unit}`")

Output shape: (262144,) Output unit: kilogram / second

Both backends return a unit-carrying tensor. Before comparing performance, express both results in kg/s and check their numerical values:

np.testing.assert_allclose(
    np.asarray(python_result.convert_to(kilogram / second).values),
    np.asarray(jax_result.convert_to(kilogram / second).values),
    rtol=1e-4,
    atol=1e-5,
)

print("Python and JAX results agree.")

Python and JAX results agree.

Small floating-point differences are expected because the two backends may evaluate equivalent numerical expressions differently.

Compare Python, cold JAX, and warm JAX¶

A fair JAX comparison should separate compilation from reuse.

For the cold measurement below, the persistent JAX compilation cache is disabled and JAX's in-memory caches are cleared immediately before evaluation. The following warm measurements then reuse the compiled calculation.

The benchmark materialises result.values as a NumPy array before stopping the clock. This ensures the returned result is complete without assuming that Tensor.values exposes a JAX array. The timed evaluation includes Bilt input validation, unit conversion, and output wrapping as well as backend execution.

import statistics
import time

import jax

from bilt import PYTHON_BACKEND

repeats = 5

python_times = []
for _ in range(repeats):
    model = CruiseImpact(**inputs)

    start = time.perf_counter()
    with PYTHON_BACKEND, cruise_impact:
        result = model.total_climate_effect
    np.asarray(result.values)
    python_times.append(time.perf_counter() - start)

python_seconds = statistics.median(python_times)

persistent_cache_enabled = jax.config.jax_enable_compilation_cache
jax.config.update("jax_enable_compilation_cache", False)

try:
    jax.clear_caches()

    model = CruiseImpact(**inputs)

    start = time.perf_counter()
    with JAX_BACKEND, cruise_impact:
        result = model.total_climate_effect
    np.asarray(result.values)
    jax_cold_seconds = time.perf_counter() - start

    jax_warm_times = []
    for _ in range(repeats):
        model = CruiseImpact(**inputs)

        start = time.perf_counter()
        with JAX_BACKEND, cruise_impact:
            result = model.total_climate_effect
        np.asarray(result.values)
        jax_warm_times.append(time.perf_counter() - start)

    jax_warm_seconds = statistics.median(jax_warm_times)
finally:
    jax.config.update(
        "jax_enable_compilation_cache",
        persistent_cache_enabled,
    )

print("| Execution | Time |")
print("| --- | ---: |")
print(f"| Python | {python_seconds * 1e3:.2f} ms |")
print(f"| JAX cold | {jax_cold_seconds * 1e3:.2f} ms |")
print(f"| JAX warm | {jax_warm_seconds * 1e3:.2f} ms |")
print()
print(f"**Warm JAX speed-up: {python_seconds / jax_warm_seconds:.2f}×**")

Execution	Time
Python	10.00 ms
JAX cold	75.99 ms
JAX warm	2.13 ms

Warm JAX speed-up: 4.70×

The exact numbers depend on the machine, JAX version, array size, and available hardware. The useful comparison is between the initial cost of compiling the calculation and the cost of running it again once compiled.

Try different ensemble sizes¶

The point at which JAX becomes worthwhile depends on the size of the calculation. The following compares several ensemble sizes.

print("| Scenarios | Python | JAX warm | Speed-up |")
print("| ---: | ---: | ---: | ---: |")

for exponent in (12, 14, 16, 18):
    scenario_count = 2**exponent
    rng = np.random.default_rng(42)

    size_inputs = {
        "altitude": tensor(
            rng.uniform(8_000.0, 12_000.0, scenario_count).astype(np.float32),
            over=scenario,
            unit=meter,
        ),
        "mach": tensor(
            rng.uniform(0.70, 0.88, scenario_count).astype(np.float32),
            over=scenario,
            unit=UNITLESS,
        ),
        "mass": tensor(
            rng.uniform(50_000.0, 80_000.0, scenario_count).astype(np.float32),
            over=scenario,
            unit=kilogram,
        ),
        "temperature_offset": tensor(
            rng.normal(0.0, 8.0, scenario_count).astype(np.float32),
            over=scenario,
            unit=kelvin,
        ),
        "wing_area": Quantity(float(np.float32(122.6)), meter**2),
        "zero_lift_drag_coefficient": Quantity(float(np.float32(0.021)), UNITLESS),
        "induced_drag_factor": Quantity(float(np.float32(0.045)), UNITLESS),
        "propulsive_efficiency": Quantity(float(np.float32(0.38)), UNITLESS),
        "nox_emission_index": Quantity(float(np.float32(0.015)), UNITLESS),
        "soot_emission_index": Quantity(float(np.float32(0.00003)), UNITLESS),
    }

    python_times = []
    for _ in range(3):
        model = CruiseImpact(**size_inputs)

        start = time.perf_counter()
        with PYTHON_BACKEND, cruise_impact:
            result = model.total_climate_effect
        python_times.append(time.perf_counter() - start)

    python_seconds = statistics.median(python_times)

    jax.clear_caches()

    model = CruiseImpact(**size_inputs)
    with JAX_BACKEND, cruise_impact:
        result = model.total_climate_effect
    np.asarray(result.values)

    jax_times = []
    for _ in range(3):
        model = CruiseImpact(**size_inputs)

        start = time.perf_counter()
        with JAX_BACKEND, cruise_impact:
            result = model.total_climate_effect
        np.asarray(result.values)
        jax_times.append(time.perf_counter() - start)

    jax_seconds = statistics.median(jax_times)

    print(
        f"| {scenario_count:,} "
        f"| {python_seconds * 1e3:.2f} ms "
        f"| {jax_seconds * 1e3:.2f} ms "
        f"| {python_seconds / jax_seconds:.2f}× |"
    )

Scenarios	Python	JAX warm	Speed-up
4,096	0.44 ms	0.43 ms	1.04×
16,384	0.83 ms	0.55 ms	1.51×
65,536	2.39 ms	0.99 ms	2.42×
262,144	11.54 ms	2.07 ms	5.59×

For small arrays, the Python backend may be entirely adequate. As the numerical work grows, reusing a compiled JAX calculation can become increasingly useful.

Current workaround: a nested JAX subproblem¶

The preferred approach is the one used above: run the complete compatible Bilt model with the JAX backend. That keeps the full calculation visible to Bilt.

There are cases where part of a model cannot currently run under JAX. As a workaround, a Python-backed transform can explicitly evaluate a separate Bilt implementation using the JAX backend.

The following handoff demonstration uses plain mathematical values rather than physical quantities. Its plain values are separate from the unit-aware cruise model above. The JAX-compatible subproblem is an ordinary Bilt model:

@system
class JaxSubproblem(metaclass=Interface):
    value: float = field()
    result: float = field()


@transform
def result(value: float) -> float:
    return value**2 + 2.0 * value + 1.0


jax_subproblem = Implementation(
    JaxSubproblem,
    [result],
)

A Python transform can delegate to it:

@transform
def output_value(input_value: float) -> float:
    from bilt import JAX_BACKEND
    from jax_backend_cruise import JaxSubproblem, jax_subproblem

    value = abs(input_value)

    with JAX_BACKEND, jax_subproblem:
        return JaxSubproblem(value=value).result

The caller still evaluates only the outer model:

from bilt import PYTHON_BACKEND

model = HybridModel(input_value=-3.0)

with PYTHON_BACKEND, hybrid_model:
    result = model.output_value

print(f"Hybrid result: {result}")

Hybrid result: 16.0

For -3.0, the Python part produces 3.0, and the nested JAX calculation returns 16.0.

This proves that the handoff is possible, but it comes with an important trade-off: the outer Bilt evaluation cannot inspect how the nested JAX result was produced. The nested calculation is therefore an opaque boundary in the otherwise transparent chain from inputs to outputs.

Use this pattern judiciously. It is best treated as an escape hatch for an unavoidable Python/JAX boundary, not as the preferred way to structure a Bilt model. Future Bilt support may provide a way to cross backend boundaries while preserving transparency across the complete calculation.

What this example shows¶

The same Bilt model can be evaluated with Python or JAX without rewriting the model.

For small or one-off calculations, Python may be the better choice. For larger calculations that are evaluated repeatedly, the cost of JAX compilation can be recovered through faster warm execution.

Where possible, select one backend around the complete Bilt model so that the full calculation remains visible and explainable. Nested backend evaluation can bridge an unavoidable compatibility boundary, but should currently be treated as a workaround.