bridgescaler.distributed_tensor

bridgescaler.distributed_tensor#

Attributes#

`CENTROID_DTYPE`
`_BATCHED_CACHE`

Classes#

`DBaseScalerTensor`	Base distributed scaler class for torch.Tensor. Used only to store attributes and methods
`DStandardScalerTensor`	Distributed version of StandardScaler. You can calculate this map-reduce style by running it on individual
`DMinMaxScalerTensor`	Distributed MinMaxScaler enables calculation of min and max of variables in datasets in parallel, then combining
`DQuantileScalerTensor`	Distributed Quantile Scaler for tensors that uses the crick TDigest Cython library to compute quantiles across multiple

Functions#

`_bucketize`(sorted_seq, values, side)	Return insertion indices of `values` into the ascending 1-D `sorted_seq` (like `torch.searchsorted`).
`fit_variable_tensor`(var_index, xv[, compression, ...])
`transform_variable_tensor`(cent_mean, cent_weight, ...)	Per-variable forward quantile transform, written branchlessly so it can be wrapped in `torch.vmap`.
`inv_transform_variable_tensor`(cent_mean, cent_weight, ...)	Per-variable inverse quantile transform, branchless for `torch.vmap`. See `transform_variable_tensor`.
`tdigest_cdf_tensor`(xv, cent_mean, cent_weight, t_min, ...)	Evaluate the TDigest CDF at every element of `xv` for a single variable.
`tdigest_quantile_tensor`(qv, cent_mean, cent_weight, ...)	Evaluate the TDigest inverse CDF (quantile function) at every element of `qv` for a single variable.
`compute_pchip_slopes`(x, y)	Vectorized Fritsch-Carlson monotone-cubic slopes (matching scipy's PchipInterpolator).
`pchip_eval_tensor`(x_knots, z_knots, m_knots, xv)	Per-variable monotone cubic-Hermite (PCHIP) evaluation of the fitted transform (fast path).

Module Contents#

bridgescaler.distributed_tensor.CENTROID_DTYPE#

bridgescaler.distributed_tensor._BATCHED_CACHE#

bridgescaler.distributed_tensor._bucketize(sorted_seq, values, side)#

Return insertion indices of values into the ascending 1-D sorted_seq (like torch.searchsorted).

torch.searchsorted has a pathologically slow Metal kernel, so on MPS the indices are computed with a broadcast reduction instead (O(N*K) memory, but roughly an order of magnitude faster there). On CPU/CUDA the binary-search torch.searchsorted is used, which is far faster and memory-light where it is well implemented. The device check is a static branch, so it does not break the graph under torch.compile.

Parameters:

sorted_seq (torch.Tensor) – ascending boundaries, shape (K,) (per variable under vmap).
values (torch.Tensor) – values to locate, any shape.
side (str) – "left" (count of boundaries strictly less than each value) or "right" (count of boundaries less than or equal to each value), matching torch.searchsorted semantics.

Returns:

int64 indices with the same shape as values.

Return type:

torch.Tensor

class bridgescaler.distributed_tensor.DBaseScalerTensor(channels_last=True)#

Base distributed scaler class for torch.Tensor. Used only to store attributes and methods shared across all distributed scaler subclasses.

x_columns_ = None#

_fit = False#

channels_last = True#

is_fit()#

extract_x_columns(x, channels_last=True)#

Extract the variable names from input x.

The variable names are expected to be stored in the variable_names attribute of the torch.Tensor. If the attribute is missing, a warning is issued to notify the user that alignment validation will be limited.

Parameters:

x (torch.Tensor) – The input tensor containing data and optionally the variable_names attribute.
channels_last (bool) – If True, then assume the variable or channel dimension is the last dimension of the array. If False, then assume the variable or channel dimension is second.

Returns:

Variable names if available; otherwise,: integer indices generated based on the length of the variable/channel dimension.

Return type:

x_columns (list[str] | list[int])

Raises:

TypeError – If x is not a torch.Tensor or if variable_names is not a list.
ValueError – If variable_names contains duplicate entries.

static extract_array(x)#

get_column_order(x_in_columns)#

Get the indices of the scaler columns that have the same name as the variables (columns) in the input x tensor. This enables users to pass a torch.Tensor to transform or inverse_transform with fewer variables than the original scaler or variables in a different order and still have the input dataset be transformed properly.

Parameters:: x_in_columns (list) – list of input variable names.
Returns:: integer indices of the input variables from x in the scaler in order.
Return type:: x_in_col_indices (list)

static package_transformed_x(x_transformed, x)#

Repackaged a transformed torch.Tensor into the same datatype as the original x, including all metadata.

Parameters:

x_transformed (torch.Tensor) – array after being transformed or inverse transformed
x (torch.Tensor) – original data

Returns:

set_channel_dim(channels_last=None)#

process_x_for_transform(x, channels_last=None)#

fit(x, weight=None)#

transform(x, channels_last=None)#

fit_transform(x, channels_last=None, weight=None)#

inverse_transform(x, channels_last=None)#

__add__(other)#

subset_columns(sel_columns)#

add_variables(other)#

static reshape_to_channels_first(stat, target)#: Reshapes ‘stat’ to align with the channel dimension (index 1).

static reshape_to_channels_last(stat, target)#: Reshapes ‘stat’ to align with the last dimension.

class bridgescaler.distributed_tensor.DStandardScalerTensor(channels_last=True)#

Bases: DBaseScalerTensor

Distributed version of StandardScaler. You can calculate this map-reduce style by running it on individual data files, returning the fitted objects, and then summing them together to represent the full dataset. Scaler supports torch.Tensor and returns a transformed tensor.

mean_x_ = None#

n_ = 0#

var_x_ = None#

fit(x, weight=None)#

transform(x, channels_last=None)#

Transform the input data from its original form to standard scaled form. If your input data has a different dimension order than the data used to fit the scaler, use the channels_last keyword argument to specify whether the new data are channels_last (True) or channels_first (False).

Parameters:

x (torch.Tensor) – Input data.
channels_last – Override the default channels_last parameter of the scaler.

Returns:

Transformed data in the same shape and type as x.

Return type:

x_transformed (torch.Tensor)

inverse_transform(x, channels_last=None)#

get_scales(x_col_order=slice(None))#

__add__(other)#

class bridgescaler.distributed_tensor.DMinMaxScalerTensor(channels_last=True)#

Bases: DBaseScalerTensor

Distributed MinMaxScaler enables calculation of min and max of variables in datasets in parallel, then combining the mins and maxes as a reduction step. Scaler supports torch.Tensor and will return a transformed tensor in the same form as the original with variable/column names preserved.

max_x_ = None#

min_x_ = None#

fit(x, weight=None)#

transform(x, channels_last=None)#

inverse_transform(x, channels_last=None)#

get_scales(x_col_order=slice(None))#

__add__(other)#

bridgescaler.distributed_tensor.fit_variable_tensor(var_index, xv, compression=None, channels_last=None)#

bridgescaler.distributed_tensor.transform_variable_tensor(cent_mean, cent_weight, t_min, t_max, n_real, xv, min_val=1e-06, max_val=0.9999999, distribution='normal')#

Per-variable forward quantile transform, written branchlessly so it can be wrapped in torch.vmap.

cent_mean/cent_weight are the padded (K,) centroid tensors for a single variable (padding means are +inf and padding weights are 0); n_real is the number of valid (non-padding) centroids.

bridgescaler.distributed_tensor.inv_transform_variable_tensor(cent_mean, cent_weight, t_min, t_max, n_real, xv, distribution='normal')#: Per-variable inverse quantile transform, branchless for torch.vmap. See transform_variable_tensor.

bridgescaler.distributed_tensor.tdigest_cdf_tensor(xv, cent_mean, cent_weight, t_min, t_max, n_real)#

Evaluate the TDigest CDF at every element of xv for a single variable.

This is a branchless, fixed-shape reformulation of the original masked implementation so that it composes with torch.vmap over the channel dimension. All four interpolation cases are computed everywhere and selected with torch.where rather than boolean-mask assignment, and centroid arrays are padded to a common length K (+inf means, 0 weights) so that n_real (not K) delimits the real centroids.

Parameters:

xv (torch.Tensor) – values to evaluate, any shape.
cent_mean (torch.Tensor) – padded (K,) centroid means, ascending, padding = +inf.
cent_weight (torch.Tensor) – padded (K,) centroid weights, padding = 0.
t_min – scalar min/max of the digest.
t_max – scalar min/max of the digest.
n_real – number of valid centroids (0-d integer tensor works under vmap).

Returns:

CDF values in [0, 1] with the same shape as xv.

Return type:

torch.Tensor

bridgescaler.distributed_tensor.tdigest_quantile_tensor(qv, cent_mean, cent_weight, t_min, t_max, n_real)#

Evaluate the TDigest inverse CDF (quantile function) at every element of qv for a single variable.

Branchless, fixed-shape reformulation of the original masked implementation for torch.vmap compatibility. See tdigest_cdf_tensor for the padding/n_real convention.

Parameters:

qv (torch.Tensor) – quantiles in [0, 1], any shape.
cent_mean (torch.Tensor) – padded (K,) centroid means, ascending, padding = +inf.
cent_weight (torch.Tensor) – padded (K,) centroid weights, padding = 0.
t_min – scalar min/max of the digest.
t_max – scalar min/max of the digest.
n_real – number of valid centroids.

Returns:

interpolated values with the same shape as qv.

Return type:

torch.Tensor

bridgescaler.distributed_tensor.compute_pchip_slopes(x, y)#

Vectorized Fritsch-Carlson monotone-cubic slopes (matching scipy’s PchipInterpolator).

Computes the derivative at every knot for a shape-preserving monotone cubic Hermite spline, for each variable (row) of x/y at once. Interior slopes use the weighted-harmonic-mean rule (set to 0 at local extrema so the spline never overshoots); endpoints use scipy’s non-centered three-point edge formula with the same monotonicity limiting.

Parameters:

x (numpy.ndarray) – strictly-increasing knot locations, shape (n_vars, M).
y (numpy.ndarray) – knot values, shape (n_vars, M).

Returns:

knot slopes, shape (n_vars, M).

Return type:

numpy.ndarray

bridgescaler.distributed_tensor.pchip_eval_tensor(x_knots, z_knots, m_knots, xv)#

Per-variable monotone cubic-Hermite (PCHIP) evaluation of the fitted transform (fast path).

Written branchlessly / fixed-shape so it composes with torch.vmap over the channel dimension. Locates each element of xv among the ascending x_knots (via _bucketize for MPS friendliness), then evaluates the Hermite cubic using the precomputed knot slopes m_knots. Values outside [x_knots[0], x_knots[-1]] saturate to the end z-knots, which equal dist_ppf(min_val) / dist_ppf(max_val) – matching the exact path’s min_val/max_val clamping.

Parameters:

x_knots (torch.Tensor) – ascending knot locations for one variable, shape (M,).
z_knots (torch.Tensor) – transformed value at each knot, shape (M,) (shared across vars).
m_knots (torch.Tensor) – PCHIP slope at each knot for one variable, shape (M,).
xv (torch.Tensor) – values to transform, any shape.

Returns:

transformed values with the same shape as xv.

Return type:

torch.Tensor

class bridgescaler.distributed_tensor.DQuantileScalerTensor(compression=250, distribution='uniform', min_val=1e-07, max_val=0.9999999, channels_last=True, compile=False, fast_transform=False, n_knots=256)#

Bases: DBaseScalerTensor

Distributed Quantile Scaler for tensors that uses the crick TDigest Cython library to compute quantiles across multiple datasets in parallel. The library can perform fitting, transforms, and inverse transforms.

DQuantileScaler supports

compression#: Recommended number of centroids to use.

distribution#: “uniform”, “normal”, or “logistic”.

min_val#: Minimum value for quantile to prevent -inf results when distribution is normal or logistic.

max_val#: Maximum value for quantile to prevent inf results when distribution is normal or logistic.

channels_last#: Whether to assume the last dim or second dim are the channel/variable dimension.

compile#: If True, transform/inverse_transform run through a cached torch.compile(fullgraph=True) of the vmapped per-variable kernel, which fuses the many elementwise ops for a sizable speedup on CPU/CUDA. Compilation is skipped on MPS (its inductor backend is immature and the _bucketize fallback already handles MPS); the default False preserves the plain eager-vmap behavior.

fast_transform#: If True, transform and inverse_transform use a monotone cubic (PCHIP) approximation of the fitted mapping instead of the exact TDigest CDF/quantile evaluation. Knots are placed at n_knots levels spaced uniformly in the output space; each direction is interpolated with a single searchsorted over n_knots (far fewer ops than the full kernel) for roughly a 3x speedup, at the cost of a small approximation error that is typically well below the TDigest’s own error (see scripts/quantile_lut_experiment.py). Default False.

n_knots#: Number of PCHIP knots per variable when fast_transform is enabled (default 256). Ignored otherwise. Must be >= 4.

compression = 250#

distribution = 'uniform'#

min_val = 1e-07#

max_val = 0.9999999#

compile = False#

fast_transform = False#

n_knots = 256#

centroids_ = None#

size_ = None#

min_ = None#

max_ = None#

centroids_mean_tensor = None#

centroids_weight_tensor = None#

min_tensor = None#

max_tensor = None#

centroids_mean_stacked = None#

centroids_weight_stacked = None#

centroids_count = None#

knots_x_ = None#

knots_z_ = None#

knots_m_ = None#

knots_minv_ = None#

td_objs_to_attributes(td_objs)#

attributes_to_td_objs()#

tensorize_attributes()#

build_stacked_centroids()#

Pad the ragged per-variable centroid lists into rectangular tensors for torch.vmap.

Each variable’s TDigest has a different number of centroids, so the per-variable (k_i,) mean/weight tensors cannot be batched directly. This pads every variable to K = max(k_i) columns, filling padding means with +inf (so they sort past all real centroids and are never selected by searchsorted) and padding weights with 0 (so they contribute nothing to the cumulative/total weights). centroids_count records the real centroid count per variable so the transform kernels know where the padding begins.

ensure_stacked_centroids()#: Rebuild the stacked centroid tensors if missing (e.g. loaded from an older serialized scaler).

build_pchip_knots()#

Build per-variable monotone-cubic (PCHIP) knots approximating the exact transform (fast path).

Knots are spaced uniformly in the output z over [dist_ppf(min_val), dist_ppf(max_val)] rather than in probability, which bounds the per-bin output error (each bin spans a fixed z-height) regardless of the input distribution – equal-probability spacing badly under-resolves the steep tails. Each z-knot is mapped to a probability via the distribution CDF, and the x-knot is the digest quantile at that probability (via tdigest_quantile_tensor). Together with the PCHIP slopes these let transform interpolate x -> z with a single searchsorted over n_knots, instead of the full TDigest CDF evaluation. Cheap and one-time; rebuilt lazily and invalidated on every fit/merge.

ensure_pchip_knots()#: Build the fast_transform knots if missing (lazily, and after load / refit / merge).

fit(x, weight=None)#

_gather_scales(x_col_order, xv)#

Select and stack the padded centroid parameters for the requested variables onto xv’s device.

Returns per-variable tensors aligned with the channel order of xv (indexed by x_col_order): stacked means (n_sel, K), stacked weights (n_sel, K), min (n_sel,), max (n_sel,), and real-centroid counts (n_sel,) — ready to be mapped over dim 0 by torch.vmap.

_get_batched_fn(kind, channel_dim, device_type)#

Build (and cache) the vmapped per-variable transform, optionally wrapped in torch.compile.

The per-variable kernel is mapped over the channel dimension (stacked params on dim 0, xv on channel_dim). When self.compile is set and the data is not on MPS, the vmapped function is wrapped in torch.compile(fullgraph=True) so inductor fuses the many elementwise ops into a few kernels. The cache is keyed by (kind, channel_dim, use_compile) so channels-first/last and compiled/eager variants coexist, and lives in a module-level weak map so nothing non-serializable lands on the scaler.

_get_fast_fn(kind, channel_dim, device_type)#

Build (and cache) the vmapped PCHIP evaluator for the fast_transform path.

Mirrors _get_batched_fn but maps pchip_eval_tensor over the channel dimension. The search/output roles swap between directions: for "transform" the search knots are the per-variable x-knots and the outputs are the shared z-knots; for "inverse" the search knots are the shared z-knots and the outputs are the per-variable x-knots. in_dims places the per-variable arg on dim 0 and the shared arg at None accordingly. Optionally torch.compiled off MPS.

transform(x, channels_last=None)#

fit_transform(x, channels_last=None, weight=None)#

inverse_transform(x, channels_last=None)#

__add__(other)#