View source on GitHub
|
Backend Abstraction Layer for Meridian.
Modules
config module: Backend configuration for Meridian.
Classes
class ExtensionType: A JAX-compatible stand-in for tf.experimental.ExtensionType.
class RNGHandler: JAX implementation of the RNGHandler using explicit key splitting.
class Tensor: Array base class for JAX
Functions
TensorShape(...): JAX implementation for TensorShape.
adstock_process(...): Dispatches adstock_process to the appropriate implementation.
computation_backend(...): Returns the active computation backend determined by inspecting the Tensor class.
computation_precision(...): Returns the active computation precision.
make_ndarray(...): JAX implementation for make_ndarray.
make_tensor_proto(...): JAX implementation for make_tensor_proto.
result_type(...): Infers the result dtype from a list of input types, backend-agnostically.
stabilize_rf_roi_grid(...): Stabilizes the RF ROI grid for JAX using a stable index lookup.
standardize_dtype(...): Converts a backend-specific dtype to a standard string representation.
to_tensor(...): Converts input data to the currently active backend tensor type.
vectorized_map(...): JAX implementation for vectorized_map.
xla_windowed_adaptive_nuts(...): JAX-specific JIT wrapper for the NUTS sampler.