require_concrete_masks

t3toolbox.backend.common.require_concrete_masks(*masks)

def require_concrete_masks(
        *masks: NDArray,  # HOST bool, static -- the uniform structure masks (must NOT be traced)
) -> None:

Guard the uniform-mask contract: masks are concrete host (numpy) arrays, never jax tracers.

Under jit any jnp op on a mask returns a tracer, which breaks the layer two ways: host-int shape/rank extraction (int(mask.sum())) raises ConcretizationTypeError, and recomputed masks leak as tracers into the (identity-hashed, never-inspected) output aux_data – silently invalid. So a traced mask here means the masks were passed among the traced jit args; the fix is functional, not numerical (raise, per the structural-vs-numerical philosophy). See docs/contributor/uniform_pytree_composition.md.

Infrastructure, not part of the ut3_* family (docs/naming_conventions.md): it guards the mask representation itself, so it serves every uniform object – plain, frame, variations, and weights.

Parameters:

masks (NDArray)

Return type:

None