pyextal.jaxGeo.objective

The epoch loss as a smooth function of the geometry, and how to tile it.

With the beam partition (ig) frozen, the loss is differentiable end to end: geometry -> tilt grid -> sg/bigkn -> structure matrix -> expm -> Voigt -> chi^2 (see jaxGeo.geometry). GeometryObjective is that function, evaluated chunk by chunk with its value, gradient and Hessian jitted.

Memory and tiling

Differentiating with respect to sg is heavier than with respect to the structure factors, because the gradient runs through the Bethe-perturbation sum over every weak beam. Measured on tilted SrTiO3 data (565 beams): 3.3 MB per tilt for a forward pass, 10.2 MB for the gradient, 17.7 MB for the exact Hessian.

XLA allocates a program’s scratch as a single block. Under WSL, JAX’s default (BFC) and ‘platform’ allocators cannot get a block over 4 GiB from CUDA however much of the card is free, which is why importing jaxGeo selects the stream-ordered allocator (cuda_async) unless told otherwise.

The geometry objective still gets its own tiling – the largest, starting from the training tile, whose derivative program fits under a block limit, read from the compiled program (fit_geometry_tile) – because a smaller card or another allocator may need it. The Voigt margin (kernel_radius each side) is kept, so small tiles waste pixels: tile 4 is 12x12 chunks of which 16 pixels count, tile 16 is 24x24 with 256. Measured, one Hessian pass on the SrTiO3 data with the previous all-tilts-at-once forward model:

tile 16, exact (8.5 GB scratch) 13 s <- chosen under cuda_async tile 16, fd (5.8 GB) 30 s tile 8, fd (2.6 GB) 13-50 s (allocator-dependent) tile 4, exact (2.6 GB) 34 s

Since jaxEngine.model.batched_forward runs a chunk’s tilts in rematerialised sub-batches, the scratch no longer grows with the tile: the exact Hessian of the whole 160x160 scan (507 beams, 99 diagonalised) as one chunk is 5.3 GB, so the search usually stops at the training tile.

hessian='auto' takes the exact Hessian where it fits and differences of the gradient otherwise. At equal tiling the objective and the training loss are the same function.

Classes

GeometryObjective

The epoch loss as a smooth function of the geometry, partition frozen.

Functions

geometry_config(config, tile)

The configuration with the scan tiled into tile-pixel chunks.

fit_geometry_tile(make_obj, config, start, mode, max_block)

Chooses the geometry tile and Hessian method that fit.

default_max_block()

Largest single scratch block to plan for, bytes, given the allocator.

Module Contents

class pyextal.jaxGeo.objective.GeometryObjective(params, metadata, config, constants, data, fns, x_ref, active, units)

The epoch loss as a smooth function of the geometry, partition frozen.

The variable is z: the offset from the reference geometry x_ref (the one the metadata was built for), in units of units, over the active parameters only. sg and bigkn enter as the stored values plus the change the JAX formulas predict, so at z = 0 the loss is exactly the stored one and float32 rounding in the formulas cannot offset it.

Parameters:

  • params (dict) – Current normalised parameters, held fixed.

  • metadata (dict) – Metadata built for x_ref.

  • config (OptConfig) – The configuration (its tiling is the one used).

  • constants (dict) – Bloch constants, with ndiag matching metadata.

  • data (dict) – Holds experimental_data and variance.

  • fns (tuple) – From GeometryModel.jax_fns.

  • x_ref (np.ndarray) – (4,) reference geometry.

  • active (np.ndarray) – Indices into x of the refined parameters.

  • units (np.ndarray) – (4,) size of one unit of each parameter.

x_ref
active
units
evals
chunks = []
_value
_vg
_vgh
property k
scratch_bytes(mode)

XLA’s scratch allocation for the derivative program, largest chunk.

XLA allocates a program’s temporaries as one block, so this is the single allocation the device must grant.

Parameters:

mode (str) – 'exact' measures the Hessian program, 'fd' the gradient one.

Returns:

Bytes.

Return type:

int

x_of(z)

Full geometry vector for offset z.

_chunk_mean(fn, z)

Mean of fn(z, chunk) over the chunks, one host transfer.

The per-chunk results are summed on the device, so the chunks are dispatched back to back instead of each waiting for the previous one’s result to reach the host.

Parameters:

  • fn (callable) – A jitted fn(z, arrays, shape); any pytree output.

  • z (np.ndarray) – (k,) offset.

Returns:

Same pytree as fn, as numpy float64.

value(z)

Mean loss over the chunks at z.

value_grad(z)

Mean loss and gradient at z.

value_grad_hess(z, mode='exact', fd_step=0.02)

Mean loss, gradient and Hessian at z.

Parameters:

  • z (np.ndarray) – (k,) offset.

  • mode (str) – 'exact' for forward-over-reverse HVPs, 'fd' for central differences of the gradient.

  • fd_step (float) – Difference step, in units, for 'fd'.

Returns:

(loss, grad (k,), hess (k, k)), Hessian symmetrised.

Return type:

tuple

pyextal.jaxGeo.objective.geometry_config(config, tile)

The configuration with the scan tiled into tile-pixel chunks.

The Voigt margin is kept.

Parameters:

  • config (OptConfig) – The training configuration.

  • tile (int) – Useful pixels per chunk side.

Returns:

Same configuration, re-tiled.

Return type:

OptConfig

pyextal.jaxGeo.objective.fit_geometry_tile(make_obj, config, start, mode, max_block)

Chooses the geometry tile and Hessian method that fit.

The limit is on the single scratch block XLA asks for, read from the compiled program, so nothing is run to find out. With mode='auto' each tile tries the exact Hessian first and then differences of the gradient: at equal tile the exact one is faster (one gradient plus k HVPs against 2k + 1 gradients), and a larger tile beats either at a smaller one, because the Voigt margin wastes a smaller share of each chunk.

Parameters:

  • make_obj (callable) – make_obj(geo_config) -> GeometryObjective.

  • config (OptConfig) – The training configuration.

  • start (int) – First tile to try; halved until something fits.

  • mode (str) – 'auto', 'exact' or 'fd'.

  • max_block (float) – Largest scratch block to allow, bytes.

Returns:

(tile, mode), mode resolved to 'exact' or 'fd'.

Return type:

tuple

Raises:

RuntimeError – If nothing fits even at tile 2.

pyextal.jaxGeo.objective.default_max_block()

Largest single scratch block to plan for, bytes, given the allocator.

Returns:

3.5 GB under the BFC/platform allocators (the WSL single-block limit is 4 GiB); otherwise 80 % of the device memory. On a large matrix that can still be too much once the training state is counted – under WSL an overflow spills into shared system memory rather than failing – so pass geo_max_block for such datasets.

Return type:

float