Skip to content

quaxed.lax¤

quaxed.lax ¤

Quaxified jax.lax.

This module wraps the functions in jax.lax with quax.quaxify. The wrapping happens dynamically through a module-level __dir__ and __getattr__. The list of available functions is in __all__ and documented in the jax.lax library.

In addition the following modules are supported:

  • quaxed.lax.linalg

The contents of these modules are likewise dynamically wrapped with quax.quaxify and their contents is listed in their respective __all__ and documented in their respective libraries.

If a function is missing, please file an Issue.

__all__ = ['abs', 'acos', 'acosh', 'add', 'approx_max_k', 'approx_min_k', 'argmax', 'argmin', 'asin', 'asinh', 'atan', 'atan2', 'atanh', 'batch_matmul', 'bessel_i0e', 'bessel_i1e', 'betainc', 'bitcast_convert_type', 'bitwise_and', 'bitwise_not', 'bitwise_or', 'bitwise_xor', 'population_count', 'broadcast', 'broadcast_in_dim', 'broadcast_shapes', 'broadcast_to_rank', 'broadcasted_iota', 'cbrt', 'ceil', 'clamp', 'clz', 'collapse', 'complex', 'concatenate', 'conj', 'conv', 'convert_element_type', 'conv_dimension_numbers', 'conv_general_dilated', 'conv_general_dilated_local', 'conv_general_dilated_patches', 'conv_transpose', 'conv_with_general_padding', 'cos', 'cosh', 'cumlogsumexp', 'cummax', 'cummin', 'cumprod', 'cumsum', 'digamma', 'div', 'dot', 'dot_general', 'dynamic_index_in_dim', 'dynamic_slice', 'dynamic_slice_in_dim', 'dynamic_update_index_in_dim', 'dynamic_update_slice', 'dynamic_update_slice_in_dim', 'eq', 'erf', 'erfc', 'erf_inv', 'exp', 'exp2', 'expand_dims', 'expm1', 'fft', 'floor', 'full', 'full_like', 'gather', 'ge', 'gt', 'igamma', 'igammac', 'imag', 'index_in_dim', 'index_take', 'integer_pow', 'iota', 'is_finite', 'le', 'lgamma', 'log', 'log1p', 'logistic', 'lt', 'max', 'min', 'mul', 'ne', 'neg', 'nextafter', 'pad', 'polygamma', 'population_count', 'pow', 'random_gamma_grad', 'real', 'reciprocal', 'reduce', 'reduce_precision', 'reduce_window', 'rem', 'reshape', 'rev', 'rng_bit_generator', 'rng_uniform', 'round', 'rsqrt', 'scatter', 'scatter_add', 'scatter_apply', 'scatter_max', 'scatter_min', 'scatter_mul', 'shift_left', 'shift_right_arithmetic', 'shift_right_logical', 'sign', 'sin', 'sinh', 'slice', 'slice_in_dim', 'sort', 'sort_key_val', 'sqrt', 'square', 'squeeze', 'sub', 'tan', 'tanh', 'top_k', 'transpose', 'zeros_like_array', 'zeta', 'associative_scan', 'cond', 'fori_loop', 'map', 'scan', 'select', 'select_n', 'switch', 'while_loop', 'stop_gradient', 'custom_linear_solve', 'custom_root', 'all_gather', 'all_to_all', 'psum', 'psum_scatter', 'pmax', 'pmin', 'pmean', 'ppermute', 'pshuffle', 'pswapaxes', 'axis_index', 'with_sharding_constraint', 'linalg', 'ConvDimensionNumbers', 'ConvGeneralDilatedDimensionNumbers', 'DotAlgorithm', 'GatherDimensionNumbers', 'GatherScatterMode', 'Precision', 'PrecisionLike', 'RandomAlgorithm', 'RoundingMethod', 'ScatterDimensionNumbers', 'DotAlgorithmPreset', 'FftType', 'zeros_like'] module-attribute ¤

Built-in mutable sequence.

If no argument is given, the constructor creates a new empty list. The argument must be an iterable if specified.