Index

_

__call__() (jax.stages.Compiled method) (jax.stages.Wrapped method) __init__() (jax.Array method) (jax.core.ClosedJaxpr method) (jax.core.Jaxpr method) (jax.custom_jvp method) (jax.custom_vjp method) (jax.Device method) (jax.dtypes.bfloat16 method) (jax.dtypes.prng_key method) (jax.experimental.checkify.Error method) (jax.experimental.sparse.BCOO method) (jax.experimental.sparse.BCSR method) (jax.experimental.sparse.COO method) (jax.experimental.sparse.CSC method) (jax.experimental.sparse.CSR method) (jax.experimental.sparse.JAXSparse method) (jax.extend.linear_util.WrappedFun method) (jax.numpy.bool_ method) (jax.numpy.character method) (jax.numpy.complex128 method) (jax.numpy.complex64 method) (jax.numpy.complexfloating method) (jax.numpy.dtype method) (jax.numpy.finfo method) (jax.numpy.flexible method) (jax.numpy.float16 method) (jax.numpy.float32 method) (jax.numpy.float64 method) (jax.numpy.floating method) (jax.numpy.generic method) (jax.numpy.iinfo method) (jax.numpy.inexact method) (jax.numpy.int16 method) (jax.numpy.int32 method) (jax.numpy.int64 method) (jax.numpy.int8 method) (jax.numpy.integer method) (jax.numpy.number method) (jax.numpy.object_ method) (jax.numpy.signedinteger method) (jax.numpy.ufunc method) (jax.numpy.uint16 method) (jax.numpy.uint32 method) (jax.numpy.uint64 method) (jax.numpy.uint8 method) (jax.numpy.unsignedinteger method) (jax.profiler.StepTraceAnnotation method) (jax.profiler.TraceAnnotation method) (jax.scipy.optimize.OptimizeResults method) (jax.scipy.spatial.transform.Rotation method) (jax.scipy.spatial.transform.Slerp method) (jax.scipy.stats.gaussian_kde method) (jax.ShapeDtypeStruct method) (jax.tree_util.Partial method)

A

abs() (in module jax.lax) (in module jax.numpy) absolute() (in module jax.numpy) acos() (in module jax.lax) acosh() (in module jax.lax) adagrad() (in module jax.example_libraries.optimizers) adam() (in module jax.example_libraries.optimizers) adamax() (in module jax.example_libraries.optimizers) add() (in module jax.lax) (in module jax.numpy) addressable_devices (jax.sharding.NamedSharding property) (jax.sharding.Sharding property) addressable_devices_indices_map() (jax.sharding.Sharding method) after_all() (in module jax.lax) all() (in module jax.numpy) all_checks (in module jax.experimental.checkify) all_gather() (in module jax.lax) all_leaves() (in module jax.tree_util) all_to_all() (in module jax.lax) allclose() (in module jax.numpy) amax() (in module jax.numpy) amin() (in module jax.numpy) angle() (in module jax.numpy) annotate_function() (in module jax.profiler) any() (in module jax.numpy) append() (in module jax.numpy) apply_along_axis() (in module jax.numpy) apply_over_axes() (in module jax.numpy) approx_max_k() (in module jax.lax) approx_min_k() (in module jax.lax) arange() (in module jax.numpy) arccos() (in module jax.numpy) arccosh() (in module jax.numpy) arcsin() (in module jax.numpy) arcsinh() (in module jax.numpy) arctan() (in module jax.numpy)

arctan2() (in module jax.numpy) arctanh() (in module jax.numpy) argmax() (in module jax.lax) (in module jax.numpy) argmin() (in module jax.lax) (in module jax.numpy) argpartition() (in module jax.numpy) argsort() (in module jax.numpy) argwhere() (in module jax.numpy) around() (in module jax.numpy) Array (class in jax) array() (in module jax.numpy) array_equal() (in module jax.numpy) array_equiv() (in module jax.numpy) array_repr() (in module jax.numpy) array_split() (in module jax.numpy) array_str() (in module jax.numpy) ArrayLike (in module jax.typing) as_text() (jax.stages.Compiled method) (jax.stages.Lowered method) asarray() (in module jax.numpy) asin() (in module jax.lax) asinh() (in module jax.lax) assert_equal() (in module jax.experimental.multihost_utils) associative_scan() (in module jax.lax) astype() (in module jax.numpy) at (jax.numpy.ndarray property) atan() (in module jax.lax) atan2() (in module jax.lax) atanh() (in module jax.lax) atleast_1d() (in module jax.numpy) atleast_2d() (in module jax.numpy) atleast_3d() (in module jax.numpy) automatic_checks (in module jax.experimental.checkify) average() (in module jax.numpy) AvgPool() (in module jax.example_libraries.stax) axis_index() (in module jax.lax)

B

ball() (in module jax.random) barrier_wait() (in module jax.experimental.host_callback) bartlett() (in module jax.numpy) batch_matmul() (in module jax.lax) BatchNorm() (in module jax.example_libraries.stax) BCOO (class in jax.experimental.sparse) bcoo_broadcast_in_dim() (in module jax.experimental.sparse) bcoo_concatenate() (in module jax.experimental.sparse) bcoo_dot_general() (in module jax.experimental.sparse) bcoo_dot_general_sampled() (in module jax.experimental.sparse) bcoo_dynamic_slice() (in module jax.experimental.sparse) bcoo_extract() (in module jax.experimental.sparse) bcoo_fromdense() (in module jax.experimental.sparse) bcoo_gather() (in module jax.experimental.sparse) bcoo_multiply_dense() (in module jax.experimental.sparse) bcoo_multiply_sparse() (in module jax.experimental.sparse) bcoo_reduce_sum() (in module jax.experimental.sparse) bcoo_reshape() (in module jax.experimental.sparse) bcoo_slice() (in module jax.experimental.sparse) bcoo_sort_indices() (in module jax.experimental.sparse) bcoo_squeeze() (in module jax.experimental.sparse) bcoo_sum_duplicates() (in module jax.experimental.sparse) bcoo_todense() (in module jax.experimental.sparse) bcoo_transpose() (in module jax.experimental.sparse) bcoo_update_layout() (in module jax.experimental.sparse) BCSR (class in jax.experimental.sparse) bcsr_dot_general() (in module jax.experimental.sparse) bcsr_extract() (in module jax.experimental.sparse) bcsr_fromdense() (in module jax.experimental.sparse) bcsr_todense() (in module jax.experimental.sparse) bernoulli() (in module jax.random) (in module jax.scipy.special) bessel_i0e() (in module jax.lax) bessel_i1e() (in module jax.lax) beta() (in module jax.random) (in module jax.scipy.special)

betainc() (in module jax.lax) (in module jax.scipy.special) betaln() (in module jax.scipy.special) bfloat16 (class in jax.dtypes) bicgstab() (in module jax.scipy.sparse.linalg) bincount() (in module jax.numpy) binomial() (in module jax.random) bitcast_convert_type() (in module jax.lax) bits (jax.numpy.finfo attribute) bits() (in module jax.random) bitwise_and() (in module jax.lax) (in module jax.numpy) bitwise_count() (in module jax.numpy) bitwise_not() (in module jax.lax) (in module jax.numpy) bitwise_or() (in module jax.lax) (in module jax.numpy) bitwise_xor() (in module jax.lax) (in module jax.numpy) blackman() (in module jax.numpy) block() (in module jax.numpy) block_diag() (in module jax.scipy.linalg) block_until_ready() (in module jax) bool_ (class in jax.numpy) breakpoint() (in module jax.debug) broadcast() (in module jax.lax) broadcast_arrays() (in module jax.numpy) broadcast_in_dim() (in module jax.lax) broadcast_one_to_all() (in module jax.experimental.multihost_utils) broadcast_shapes() (in module jax.lax) (in module jax.numpy) broadcast_to() (in module jax.numpy) broadcast_to_rank() (in module jax.lax) broadcasted_iota() (in module jax.lax) build_tree() (in module jax.tree_util)

C

c_ (in module jax.numpy) cache() (in module jax.extend.linear_util) call() (in module jax.experimental.host_callback) callback() (in module jax.debug) CallbackException can_cast() (in module jax.numpy) canonicalize_dtype() (in module jax.dtypes) categorical() (in module jax.random) cauchy() (in module jax.random) cbrt() (in module jax.lax) (in module jax.numpy) cdf() (in module jax.scipy.stats.bernoulli) (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.gennorm) (in module jax.scipy.stats.laplace) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.truncnorm) cdouble (in module jax.numpy) ceil() (in module jax.lax) (in module jax.numpy) celu() (in module jax.nn) cg() (in module jax.scipy.sparse.linalg) character (class in jax.numpy) check() (in module jax.experimental.checkify) check_error() (in module jax.experimental.checkify) check_tracer_leaks (in module jax) checkify() (in module jax.experimental.checkify) checking_leaks (in module jax) checkpoint() (in module jax) chisquare() (in module jax.random) cho_factor() (in module jax.scipy.linalg) cho_solve() (in module jax.scipy.linalg) choice() (in module jax.random) cholesky() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) choose() (in module jax.numpy) clamp() (in module jax.lax) clear_caches() (in module jax) clip() (in module jax.numpy) clip_grads() (in module jax.example_libraries.optimizers) ClosedJaxpr (class in jax.core) closure_convert() (in module jax) clz() (in module jax.lax) collapse() (in module jax.lax) column_stack() (in module jax.numpy) compile() (jax.stages.Lowered method) Compiled (class in jax.stages) compiler_ir() (jax.stages.Lowered method) complex() (in module jax.lax) complex128 (class in jax.numpy) complex64 (class in jax.numpy) complex_ (in module jax.numpy) complexfloating (class in jax.numpy) ComplexWarning compress() (in module jax.numpy) concatenate() (in module jax.lax) (in module jax.numpy) ConcretizationTypeError (class in jax.errors) cond() (in module jax.lax) (in module jax.numpy.linalg)

config (in module jax) conj() (in module jax.lax) (in module jax.numpy) conjugate() (in module jax.numpy) constant() (in module jax.example_libraries.optimizers) (in module jax.nn.initializers) Conv() (in module jax.example_libraries.stax) conv() (in module jax.lax) Conv1DTranspose() (in module jax.example_libraries.stax) conv_dimension_numbers() (in module jax.lax) conv_general_dilated() (in module jax.lax) conv_general_dilated_local() (in module jax.lax) conv_general_dilated_patches() (in module jax.lax) conv_transpose() (in module jax.lax) conv_with_general_padding() (in module jax.lax) ConvDimensionNumbers (class in jax.lax) convert_element_type() (in module jax.lax) ConvGeneralDilatedDimensionNumbers (in module jax.lax) convolve() (in module jax.numpy) (in module jax.scipy.signal) convolve2d() (in module jax.scipy.signal) ConvTranspose() (in module jax.example_libraries.stax) COO (class in jax.experimental.sparse) coo_fromdense() (in module jax.experimental.sparse) coo_matmat() (in module jax.experimental.sparse) coo_matvec() (in module jax.experimental.sparse) coo_todense() (in module jax.experimental.sparse) copy() (in module jax.numpy) copysign() (in module jax.numpy) corrcoef() (in module jax.numpy) correlate() (in module jax.numpy) (in module jax.scipy.signal) correlate2d() (in module jax.scipy.signal) cos() (in module jax.lax) (in module jax.numpy) cosh() (in module jax.lax) (in module jax.numpy) cost_analysis() (jax.stages.Compiled method) (jax.stages.Lowered method) count_nonzero() (in module jax.numpy) cov() (in module jax.numpy) CPU cross() (in module jax.numpy) CSC (class in jax.experimental.sparse) csd() (in module jax.scipy.signal) csingle (in module jax.numpy) CSR (class in jax.experimental.sparse) csr_fromdense() (in module jax.experimental.sparse) csr_matmat() (in module jax.experimental.sparse) csr_matvec() (in module jax.experimental.sparse) csr_todense() (in module jax.experimental.sparse) cumlogsumexp() (in module jax.lax) cummax() (in module jax.lax) cummin() (in module jax.lax) cumprod() (in module jax.lax) (in module jax.numpy) cumsum() (in module jax.lax) (in module jax.numpy) custom_gradient() (in module jax) custom_jvp (class in jax) custom_linear_solve() (in module jax.lax) custom_partitioning() (in module jax.experimental.custom_partitioning) custom_root() (in module jax.lax) custom_vjp (class in jax)

D

data (jax.experimental.sparse.BCOO attribute) dct() (in module jax.scipy.fft) dctn() (in module jax.scipy.fft) debug_infs (in module jax) debug_nans (in module jax) default() (jax.sharding.PmapSharding class method) default_backend() (in module jax) (in module jax.lib.xla_bridge) default_device (in module jax) default_matmul_precision (in module jax) default_prng_impl (in module jax) define_prng_impl() (in module jax.extend.random) deg2rad() (in module jax.numpy) degrees() (in module jax.numpy) delete() (in module jax.numpy) delta_orthogonal() (in module jax.nn.initializers) Dense() (in module jax.example_libraries.stax) det (in module jax.numpy.linalg) det() (in module jax.scipy.linalg) Device (class in jax) device_count() (in module jax) device_get() (in module jax) device_memory_profile() (in module jax.profiler) device_put() (in module jax) device_put_replicated() (in module jax) device_put_sharded() (in module jax) device_set (jax.sharding.GSPMDSharding property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.PositionalSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) DeviceArray devices() (in module jax) devices_indices_map() (jax.sharding.GSPMDSharding method) (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) (jax.sharding.SingleDeviceSharding method) (jax.sharding.XLACompatibleSharding method)

diag() (in module jax.numpy) diag_indices() (in module jax.numpy) diag_indices_from() (in module jax.numpy) diagflat() (in module jax.numpy) diagonal() (in module jax.numpy) diff() (in module jax.numpy) digamma() (in module jax.lax) (in module jax.scipy.special) digitize() (in module jax.numpy) dirichlet() (in module jax.random) disable_jit() (in module jax) disable_x64() (in module jax.experimental) div() (in module jax.lax) div_checks (in module jax.experimental.checkify) divide() (in module jax.numpy) divmod() (in module jax.numpy) dot() (in module jax.lax) (in module jax.numpy) dot_general() (in module jax.lax) double (in module jax.numpy) double_sided_maxwell() (in module jax.random) Dropout() (in module jax.example_libraries.stax) dsplit() (in module jax.numpy) dstack() (in module jax.numpy) dtype (class in jax.numpy) (jax.numpy.finfo attribute) DTypeLike (in module jax.typing) dynamic_index_in_dim() (in module jax.lax) dynamic_slice() (in module jax.lax) dynamic_slice_in_dim() (in module jax.lax) dynamic_update_index_in_dim() (in module jax.lax) dynamic_update_slice() (in module jax.lax) dynamic_update_slice_in_dim() (in module jax.lax)

E

ediff1d() (in module jax.numpy) eig() (in module jax.lax.linalg) (in module jax.numpy.linalg) eigh() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) eigh_tridiagonal() (in module jax.scipy.linalg) eigvals() (in module jax.numpy.linalg) eigvalsh() (in module jax.numpy.linalg) einsum() (in module jax.numpy) einsum_path() (in module jax.numpy) elementwise() (in module jax.example_libraries.stax) elu() (in module jax.nn) empty() (in module jax.experimental.sparse) (in module jax.numpy) empty_like() (in module jax.numpy) enable_checks (in module jax) enable_custom_prng (in module jax) enable_custom_vjp_by_custom_transpose (in module jax) enable_x64() (in module jax.experimental) ensure_compile_time_eval() (in module jax) entr() (in module jax.scipy.special) eps (jax.numpy.finfo attribute) epsneg (jax.numpy.finfo attribute) eq() (in module jax.lax) equal() (in module jax.numpy)

erf() (in module jax.lax) (in module jax.scipy.special) erf_inv() (in module jax.lax) erfc() (in module jax.lax) (in module jax.scipy.special) erfinv() (in module jax.scipy.special) Error (class in jax.experimental.checkify) eval_shape() (in module jax) evaluate() (jax.scipy.stats.gaussian_kde method) exp() (in module jax.lax) (in module jax.numpy) exp1() (in module jax.scipy.special) exp2() (in module jax.numpy) expand_dims() (in module jax.lax) (in module jax.numpy) expi (in module jax.scipy.special) expit() (in module jax.scipy.special) expm() (in module jax.scipy.linalg) expm1() (in module jax.lax) (in module jax.numpy) expm_frechet() (in module jax.scipy.linalg) expn (in module jax.scipy.special) exponential() (in module jax.random) exponential_decay() (in module jax.example_libraries.optimizers) extract() (in module jax.numpy) eye() (in module jax.experimental.sparse) (in module jax.numpy)

F

f() (in module jax.random) fabs() (in module jax.numpy) factorial() (in module jax.scipy.special) FanInConcat() (in module jax.example_libraries.stax) FanOut() (in module jax.example_libraries.stax) fft() (in module jax.lax) (in module jax.numpy.fft) fft2() (in module jax.numpy.fft) fftconvolve() (in module jax.scipy.signal) fftfreq() (in module jax.numpy.fft) fftn() (in module jax.numpy.fft) fftshift() (in module jax.numpy.fft) fill_diagonal() (in module jax.numpy) finfo (class in jax.numpy) fix() (in module jax.numpy) flatnonzero() (in module jax.numpy) flexible (class in jax.numpy) flip() (in module jax.numpy) fliplr() (in module jax.numpy) flipud() (in module jax.numpy) float0 (in module jax.dtypes) float16 (class in jax.numpy) float32 (class in jax.numpy) float64 (class in jax.numpy) float_ (in module jax.numpy) float_checks (in module jax.experimental.checkify)

float_power() (in module jax.numpy) floating (class in jax.numpy) floor() (in module jax.lax) (in module jax.numpy) floor_divide() (in module jax.numpy) fmax() (in module jax.numpy) fmin() (in module jax.numpy) fmod() (in module jax.numpy) fold_in() (in module jax.random) fori_loop() (in module jax.lax) forward-mode autodiff frexp() (in module jax.numpy) from_dlpack() (in module jax.dlpack) (in module jax.numpy) frombuffer() (in module jax.numpy) fromfile() (in module jax.numpy) fromfunction() (in module jax.numpy) fromiter() (in module jax.numpy) frompyfunc() (in module jax.numpy) fromstring() (in module jax.numpy) full() (in module jax.lax) (in module jax.numpy) full_like() (in module jax.lax) (in module jax.numpy) functional programming funm() (in module jax.scipy.linalg)

G

gamma() (in module jax.random) (in module jax.scipy.special) gammainc() (in module jax.scipy.special) gammaincc() (in module jax.scipy.special) gammaln() (in module jax.scipy.special) gather() (in module jax.lax) GatherDimensionNumbers (class in jax.lax) GatherScatterMode (class in jax.lax) gaussian_kde (class in jax.scipy.stats) gcd() (in module jax.numpy) ge() (in module jax.lax) gelu() (in module jax.nn) GeneralConv() (in module jax.example_libraries.stax) GeneralConvTranspose() (in module jax.example_libraries.stax) generalized_normal() (in module jax.random) generic (class in jax.numpy) geometric() (in module jax.random)

geomspace() (in module jax.numpy) get_backend() (in module jax.lib.xla_bridge) get_compile_options() (in module jax.lib.xla_bridge) get_printoptions() (in module jax.numpy) global_array_to_host_local_array() (in module jax.experimental.multihost_utils) glorot_normal() (in module jax.nn.initializers) glorot_uniform() (in module jax.nn.initializers) glu() (in module jax.nn) gmres() (in module jax.scipy.sparse.linalg) GPU grad() (in module jax) (in module jax.experimental.sparse) gradient() (in module jax.numpy) greater() (in module jax.numpy) greater_equal() (in module jax.numpy) GSPMDSharding (class in jax.sharding) gt() (in module jax.lax) gumbel() (in module jax.random)

H

hamming() (in module jax.numpy) hanning() (in module jax.numpy) hard_sigmoid() (in module jax.nn) hard_silu() (in module jax.nn) hard_swish() (in module jax.nn) hard_tanh() (in module jax.nn) he_normal() (in module jax.nn.initializers) he_uniform() (in module jax.nn.initializers) heaviside() (in module jax.numpy) hessenberg() (in module jax.lax.linalg) (in module jax.scipy.linalg)

hessian() (in module jax) hfft() (in module jax.numpy.fft) histogram() (in module jax.numpy) histogram2d() (in module jax.numpy) histogram_bin_edges() (in module jax.numpy) histogramdd() (in module jax.numpy) host_local_array_to_global_array() (in module jax.experimental.multihost_utils) householder_product() (in module jax.lax.linalg) hsplit() (in module jax.numpy) hstack() (in module jax.numpy) hyp1f1 (in module jax.scipy.special) hypot() (in module jax.numpy)

I

i0 (in module jax.numpy) i0() (in module jax.scipy.special) i0e() (in module jax.scipy.special) i1() (in module jax.scipy.special) i1e() (in module jax.scipy.special) id_print() (in module jax.experimental.host_callback) id_tap() (in module jax.experimental.host_callback) idct() (in module jax.scipy.fft) idctn() (in module jax.scipy.fft) identity() (in module jax.numpy) iexp (jax.numpy.finfo attribute) ifft() (in module jax.numpy.fft) ifft2() (in module jax.numpy.fft) ifftn() (in module jax.numpy.fft) ifftshift() (in module jax.numpy.fft) igamma() (in module jax.lax) igammac() (in module jax.lax) ihfft() (in module jax.numpy.fft) iinfo (class in jax.numpy) imag() (in module jax.lax) (in module jax.numpy) in1d() (in module jax.numpy) in_tree (jax.stages.Compiled property) (jax.stages.Lowered property) index_checks (in module jax.experimental.checkify) index_exp (in module jax.numpy) index_in_dim() (in module jax.lax) index_take() (in module jax.lax) indices (jax.experimental.sparse.BCOO attribute) indices() (in module jax.numpy) inexact (class in jax.numpy) init_fn (jax.example_libraries.optimizers.Optimizer attribute) initialize() (in module jax.distributed) initialize_cache() (in module jax.experimental.compilation_cache.compilation_cache) inner() (in module jax.numpy) insert() (in module jax.numpy) inspect_array_sharding() (in module jax.debug) int16 (class in jax.numpy) int32 (class in jax.numpy) int64 (class in jax.numpy) int8 (class in jax.numpy) int_ (in module jax.numpy) integer (class in jax.numpy) integer_pow() (in module jax.lax) integrate_box_1d() (jax.scipy.stats.gaussian_kde method) integrate_gaussian() (jax.scipy.stats.gaussian_kde method) integrate_kde() (jax.scipy.stats.gaussian_kde method)

interp() (in module jax.numpy) intersect1d() (in module jax.numpy) inv() (in module jax.numpy.linalg) (in module jax.scipy.linalg) inverse_time_decay() (in module jax.example_libraries.optimizers) invert() (in module jax.numpy) io_callback() (in module jax.experimental) iota() (in module jax.lax) irfft() (in module jax.numpy.fft) irfft2() (in module jax.numpy.fft) irfftn() (in module jax.numpy.fft) is_equivalent_to() (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) (jax.sharding.XLACompatibleSharding method) is_finite() (in module jax.lax) is_fully_addressable (jax.sharding.GSPMDSharding property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.PositionalSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) is_fully_replicated (jax.sharding.GSPMDSharding property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.PositionalSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) is_initialized() (in module jax.experimental.compilation_cache.compilation_cache) isclose() (in module jax.numpy) iscomplex() (in module jax.numpy) iscomplexobj() (in module jax.numpy) isf() (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) isfinite() (in module jax.numpy) isin() (in module jax.numpy) isinf() (in module jax.numpy) isnan() (in module jax.numpy) isneginf() (in module jax.numpy) isposinf() (in module jax.numpy) isreal() (in module jax.numpy) isrealobj() (in module jax.numpy) isscalar() (in module jax.numpy) issubdtype() (in module jax.dtypes) (in module jax.numpy) istft() (in module jax.scipy.signal) iterable() (in module jax.numpy) ix_() (in module jax.numpy)

J

jacfwd() (in module jax) jacrev() (in module jax) jax.core module jax.debug module jax.distributed module jax.dlpack module jax.dtypes module jax.example_libraries module jax.example_libraries.optimizers module jax.example_libraries.stax module jax.experimental.array_api module jax.experimental.checkify module jax.experimental.compilation_cache.compilation_cache module jax.experimental.custom_partitioning module jax.experimental.host_callback module jax.experimental.jet module jax.experimental.maps module jax.experimental.multihost_utils module jax.experimental.pjit module jax.experimental.sparse module jax.experimental.sparse.linalg module jax.extend module jax.extend.linear_util module jax.extend.mlir module jax.extend.mlir.dialects module jax.extend.mlir.ir module jax.extend.mlir.passmanager module jax.extend.random module jax.flatten_util module jax.image module jax.lax module jax.lax.linalg module jax.nn module jax.nn.initializers module jax.numpy module jax.numpy.fft module jax.numpy.linalg module jax.ops module jax.profiler module jax.random module jax.scipy.fft module jax.scipy.integrate module

jax.scipy.linalg module jax.scipy.ndimage module jax.scipy.optimize module jax.scipy.signal module jax.scipy.sparse.linalg module jax.scipy.spatial.transform module jax.scipy.special module jax.scipy.stats module jax.scipy.stats.bernoulli module jax.scipy.stats.beta module jax.scipy.stats.betabinom module jax.scipy.stats.binom module jax.scipy.stats.cauchy module jax.scipy.stats.chi2 module jax.scipy.stats.dirichlet module jax.scipy.stats.expon module jax.scipy.stats.gamma module jax.scipy.stats.gennorm module jax.scipy.stats.geom module jax.scipy.stats.laplace module jax.scipy.stats.logistic module jax.scipy.stats.multinomial module jax.scipy.stats.multivariate_normal module jax.scipy.stats.nbinom module jax.scipy.stats.norm module jax.scipy.stats.pareto module jax.scipy.stats.poisson module jax.scipy.stats.t module jax.scipy.stats.truncnorm module jax.scipy.stats.uniform module jax.scipy.stats.vonmises module jax.scipy.stats.wrapcauchy module jax.sharding module jax.stages module jax.tree_util module jax.typing module jaxpr Jaxpr (class in jax.core) JaxRuntimeError JAXSparse (class in jax.experimental.sparse) jet() (in module jax.experimental.jet) JIT jit() (in module jax) JoinPoint (class in jax.example_libraries.optimizers) JVP jvp() (in module jax)

K

kaiser() (in module jax.numpy) key() (in module jax.random) key_data() (in module jax.random)

keystr() (in module jax.tree_util) kl_div() (in module jax.scipy.special) kron() (in module jax.numpy)

L

l2_norm() (in module jax.example_libraries.optimizers) laplace() (in module jax.random) lcm() (in module jax.numpy) ldexp() (in module jax.numpy) le() (in module jax.lax) leaky_relu() (in module jax.nn) lecun_normal() (in module jax.nn.initializers) lecun_uniform() (in module jax.nn.initializers) left_shift() (in module jax.numpy) less() (in module jax.numpy) less_equal() (in module jax.numpy) lexsort() (in module jax.numpy) lgamma() (in module jax.lax) linear_transpose() (in module jax) linearize() (in module jax) linspace() (in module jax.numpy) live_arrays() (in module jax) load() (in module jax.numpy) lobpcg_standard() (in module jax.experimental.sparse.linalg) local_device_count() (in module jax) local_devices() (in module jax) log() (in module jax.lax) (in module jax.numpy) log10() (in module jax.numpy) log1p() (in module jax.lax) (in module jax.numpy) log2() (in module jax.numpy) log_compiles (in module jax) log_ndtr (in module jax.scipy.special) log_sigmoid() (in module jax.nn) log_softmax() (in module jax.nn) logaddexp (in module jax.numpy) logaddexp2 (in module jax.numpy) logcdf() (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.truncnorm) loggamma() (in module jax.random) logical_and() (in module jax.numpy) logical_not() (in module jax.numpy) logical_or() (in module jax.numpy) logical_xor() (in module jax.numpy) logistic() (in module jax.lax) (in module jax.random)

logit (in module jax.scipy.special) lognormal() (in module jax.random) logpdf() (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.dirichlet) (in module jax.scipy.stats.expon) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.gennorm) (in module jax.scipy.stats.laplace) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.multivariate_normal) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.pareto) (in module jax.scipy.stats.t) (in module jax.scipy.stats.truncnorm) (in module jax.scipy.stats.uniform) (in module jax.scipy.stats.vonmises) (in module jax.scipy.stats.wrapcauchy) (jax.scipy.stats.gaussian_kde method) logpmf() (in module jax.scipy.stats.bernoulli) (in module jax.scipy.stats.betabinom) (in module jax.scipy.stats.binom) (in module jax.scipy.stats.geom) (in module jax.scipy.stats.multinomial) (in module jax.scipy.stats.nbinom) (in module jax.scipy.stats.poisson) logsf() (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.truncnorm) logspace() (in module jax.numpy) logsumexp() (in module jax.nn) (in module jax.scipy.special) lower() (jax.stages.Wrapped method) Lowered (class in jax.stages) lpmn() (in module jax.scipy.special) lpmn_values() (in module jax.scipy.special) lstsq() (in module jax.numpy.linalg) lt() (in module jax.lax) lu() (in module jax.lax.linalg) (in module jax.scipy.linalg) lu_factor() (in module jax.scipy.linalg) lu_solve() (in module jax.scipy.linalg)

M

machep (jax.numpy.finfo attribute) make_array_from_callback() (in module jax) make_array_from_single_device_arrays() (in module jax) make_jaxpr() (in module jax) make_schedule() (in module jax.example_libraries.optimizers) map() (in module jax.lax) map_coordinates() (in module jax.scipy.ndimage) mask_indices() (in module jax.numpy) matmul() (in module jax.numpy) matrix_power() (in module jax.numpy.linalg) matrix_rank() (in module jax.numpy.linalg) matrix_transpose() (in module jax.numpy) max (jax.numpy.finfo attribute) max() (in module jax.lax) (in module jax.numpy) maxexp (jax.numpy.finfo attribute) maximum() (in module jax.numpy) MaxPool() (in module jax.example_libraries.stax) maxwell() (in module jax.random) mean() (in module jax.numpy) median() (in module jax.numpy) memory_analysis() (jax.stages.Compiled method) memory_kind (jax.sharding.GSPMDSharding property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.PositionalSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) merge_linear_aux() (in module jax.extend.linear_util) Mesh (class in jax.sharding) meshgrid() (in module jax.numpy) mgrid (in module jax.numpy) min (jax.numpy.finfo attribute) min() (in module jax.lax) (in module jax.numpy) minexp (jax.numpy.finfo attribute) minimize() (in module jax.scipy.optimize) minimum() (in module jax.numpy) mod() (in module jax.numpy) mode() (in module jax.scipy.stats) modf() (in module jax.numpy) module jax.core jax.debug jax.distributed jax.dlpack jax.dtypes jax.example_libraries jax.example_libraries.optimizers jax.example_libraries.stax jax.experimental.array_api jax.experimental.checkify jax.experimental.compilation_cache.compilation_cache jax.experimental.custom_partitioning jax.experimental.host_callback jax.experimental.jet jax.experimental.maps jax.experimental.multihost_utils jax.experimental.pjit jax.experimental.sparse jax.experimental.sparse.linalg jax.extend jax.extend.linear_util jax.extend.mlir jax.extend.mlir.dialects jax.extend.mlir.ir jax.extend.mlir.passmanager jax.extend.random jax.flatten_util jax.image jax.lax jax.lax.linalg jax.nn jax.nn.initializers jax.numpy jax.numpy.fft jax.numpy.linalg jax.ops jax.profiler jax.random jax.scipy.fft jax.scipy.integrate jax.scipy.linalg jax.scipy.ndimage jax.scipy.optimize jax.scipy.signal jax.scipy.sparse.linalg jax.scipy.spatial.transform jax.scipy.special jax.scipy.stats jax.scipy.stats.bernoulli jax.scipy.stats.beta jax.scipy.stats.betabinom jax.scipy.stats.binom jax.scipy.stats.cauchy jax.scipy.stats.chi2 jax.scipy.stats.dirichlet jax.scipy.stats.expon jax.scipy.stats.gamma jax.scipy.stats.gennorm jax.scipy.stats.geom jax.scipy.stats.laplace jax.scipy.stats.logistic jax.scipy.stats.multinomial jax.scipy.stats.multivariate_normal jax.scipy.stats.nbinom jax.scipy.stats.norm jax.scipy.stats.pareto jax.scipy.stats.poisson jax.scipy.stats.t jax.scipy.stats.truncnorm jax.scipy.stats.uniform jax.scipy.stats.vonmises jax.scipy.stats.wrapcauchy jax.sharding jax.stages jax.tree_util jax.typing

momentum() (in module jax.example_libraries.optimizers) moveaxis() (in module jax.numpy) mul() (in module jax.lax) multi_dot() (in module jax.numpy.linalg) multigammaln() (in module jax.scipy.special) multiply() (in module jax.numpy) multivariate_normal() (in module jax.random)

N

named_call() (in module jax) named_scope() (in module jax) NamedSharding (class in jax.sharding) nan_checks (in module jax.experimental.checkify) nan_to_num() (in module jax.numpy) nanargmax() (in module jax.numpy) nanargmin() (in module jax.numpy) nancumprod() (in module jax.numpy) nancumsum() (in module jax.numpy) nanmax() (in module jax.numpy) nanmean() (in module jax.numpy) nanmedian() (in module jax.numpy) nanmin() (in module jax.numpy) nanpercentile() (in module jax.numpy) nanprod() (in module jax.numpy) nanquantile() (in module jax.numpy) nanstd() (in module jax.numpy) nansum() (in module jax.numpy) nanvar() (in module jax.numpy) ndarray (in module jax.numpy)

ndim() (in module jax.numpy) ndtr() (in module jax.scipy.special) ndtri() (in module jax.scipy.special) ne() (in module jax.lax) neg() (in module jax.lax) negative() (in module jax.numpy) negep (jax.numpy.finfo attribute) nesterov() (in module jax.example_libraries.optimizers) nexp (jax.numpy.finfo attribute) nextafter() (in module jax.lax) (in module jax.numpy) nmant (jax.numpy.finfo attribute) NonConcreteBooleanIndexError (class in jax.errors) nonzero() (in module jax.numpy) norm() (in module jax.numpy.linalg) normal() (in module jax.nn.initializers) (in module jax.random) not_equal() (in module jax.numpy) number (class in jax.numpy) numpy_rank_promotion (in module jax)

O

object_ (class in jax.numpy) ogrid (in module jax.numpy) one_hot() (in module jax.nn) ones() (in module jax.nn.initializers) (in module jax.numpy) ones_like() (in module jax.numpy)

Optimizer (class in jax.example_libraries.optimizers) optimizer() (in module jax.example_libraries.optimizers) OptimizeResults (class in jax.scipy.optimize) OptimizerState (class in jax.example_libraries.optimizers) orthogonal() (in module jax.nn.initializers) (in module jax.random) outer() (in module jax.numpy)

P

pack_optimizer_state() (in module jax.example_libraries.optimizers) packbits() (in module jax.numpy) packed_state (jax.example_libraries.optimizers.OptimizerState attribute) pad() (in module jax.lax) (in module jax.numpy) parallel() (in module jax.example_libraries.stax) params_fn (jax.example_libraries.optimizers.Optimizer attribute) pareto() (in module jax.random) Partial (class in jax.tree_util) partition() (in module jax.numpy) PartitionSpec (class in jax.sharding) pdf() (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.dirichlet) (in module jax.scipy.stats.expon) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.gennorm) (in module jax.scipy.stats.laplace) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.multivariate_normal) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.pareto) (in module jax.scipy.stats.t) (in module jax.scipy.stats.truncnorm) (in module jax.scipy.stats.uniform) (in module jax.scipy.stats.vonmises) (in module jax.scipy.stats.wrapcauchy) (jax.scipy.stats.gaussian_kde method) pdot() (in module jax.lax) percentile() (in module jax.numpy) permutation() (in module jax.random) piecewise() (in module jax.numpy) piecewise_constant() (in module jax.example_libraries.optimizers) pinv (in module jax.numpy.linalg) pjit() (in module jax.experimental.pjit) place() (in module jax.numpy) pmap() (in module jax) PmapSharding (class in jax.sharding) pmax() (in module jax.lax) pmean() (in module jax.lax) pmf() (in module jax.scipy.stats.bernoulli) (in module jax.scipy.stats.betabinom) (in module jax.scipy.stats.binom) (in module jax.scipy.stats.geom) (in module jax.scipy.stats.multinomial) (in module jax.scipy.stats.nbinom) (in module jax.scipy.stats.poisson)

pmin() (in module jax.lax) poch (in module jax.scipy.special) poisson() (in module jax.random) polar() (in module jax.scipy.linalg) poly() (in module jax.numpy) polyadd() (in module jax.numpy) polyder() (in module jax.numpy) polydiv() (in module jax.numpy) polyfit() (in module jax.numpy) polygamma() (in module jax.lax) (in module jax.scipy.special) polyint() (in module jax.numpy) polymul() (in module jax.numpy) polynomial_decay() (in module jax.example_libraries.optimizers) polysub() (in module jax.numpy) polyval() (in module jax.numpy) population_count() (in module jax.lax) PositionalSharding (class in jax.sharding) positive() (in module jax.numpy) pow() (in module jax.lax) power() (in module jax.numpy) ppermute() (in module jax.lax) ppf() (in module jax.scipy.stats.bernoulli) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) Precision (class in jax.lax) precision (jax.numpy.finfo attribute) PrecisionLike (in module jax.lax) print() (in module jax.debug) print_environment_info() (in module jax) printoptions() (in module jax.numpy) prng_key (class in jax.dtypes) PRNGKey() (in module jax.random) process_allgather() (in module jax.experimental.multihost_utils) process_count() (in module jax) process_index() (in module jax) prod() (in module jax.numpy) promote_types() (in module jax.numpy) pshuffle() (in module jax.lax) psum() (in module jax.lax) psum_scatter() (in module jax.lax) pswapaxes() (in module jax.lax) ptp() (in module jax.numpy) pure function pure_callback() (in module jax) put() (in module jax.numpy)

Q

qdwh() (in module jax.lax.linalg) qr() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg)

quantile() (in module jax.numpy)

R

r_ (in module jax.numpy) rad2deg() (in module jax.numpy) rademacher() (in module jax.random) radians() (in module jax.numpy) randint() (in module jax.random) random_bcoo() (in module jax.experimental.sparse) random_gamma_grad() (in module jax.lax) rankdata() (in module jax.scipy.stats) ravel() (in module jax.numpy) ravel_multi_index() (in module jax.numpy) ravel_pytree() (in module jax.flatten_util) rayleigh() (in module jax.random) rbg_prng_impl (in module jax.extend.random) real() (in module jax.lax) (in module jax.numpy) reciprocal() (in module jax.lax) (in module jax.numpy) reduce() (in module jax.lax) reduce_precision() (in module jax.lax) reduce_window() (in module jax.lax) register_pytree_node() (in module jax.tree_util) register_pytree_node_class() (in module jax.tree_util) register_pytree_with_keys() (in module jax.tree_util) register_pytree_with_keys_class() (in module jax.tree_util) rel_entr() (in module jax.scipy.special) relu (in module jax.nn) relu6 (in module jax.nn) rem() (in module jax.lax) remainder() (in module jax.numpy) repeat() (in module jax.numpy) resample() (jax.scipy.stats.gaussian_kde method) reset_cache() (in module jax.experimental.compilation_cache.compilation_cache)

reshape() (in module jax.lax) (in module jax.numpy) resize() (in module jax.image) (in module jax.numpy) ResizeMethod (class in jax.image) resolution (jax.numpy.finfo attribute) result_type() (in module jax.dtypes) (in module jax.numpy) rev() (in module jax.lax) reverse-mode autodiff rfft() (in module jax.numpy.fft) rfft2() (in module jax.numpy.fft) rfftfreq() (in module jax.numpy.fft) rfftn() (in module jax.numpy.fft) right_shift() (in module jax.numpy) rint() (in module jax.numpy) rmsprop() (in module jax.example_libraries.optimizers) rmsprop_momentum() (in module jax.example_libraries.optimizers) rng_bit_generator() (in module jax.lax) rng_uniform() (in module jax.lax) roll() (in module jax.numpy) rollaxis() (in module jax.numpy) roots() (in module jax.numpy) rot90() (in module jax.numpy) Rotation (class in jax.scipy.spatial.transform) round() (in module jax.lax) (in module jax.numpy) round_() (in module jax.numpy) RoundingMethod (class in jax.lax) rsf2csf() (in module jax.scipy.linalg) rsqrt() (in module jax.lax) runtime_executable() (jax.stages.Compiled method)

S

s_ (in module jax.numpy) save() (in module jax.numpy) save_device_memory_profile() (in module jax.profiler) savez() (in module jax.numpy) scalar_type_of() (in module jax.dtypes) scale_and_translate() (in module jax.image) scan() (in module jax.lax) scatter() (in module jax.lax) scatter_add() (in module jax.lax) scatter_apply() (in module jax.lax) scatter_max() (in module jax.lax) scatter_min() (in module jax.lax) scatter_mul() (in module jax.lax) ScatterDimensionNumbers (class in jax.lax) schur() (in module jax.lax.linalg) (in module jax.scipy.linalg) searchsorted() (in module jax.numpy) seed_with_impl() (in module jax.extend.random) segment_max() (in module jax.ops) segment_min() (in module jax.ops) segment_prod() (in module jax.ops) segment_sum() (in module jax.ops) select() (in module jax.lax) (in module jax.numpy) select_n() (in module jax.lax) selu() (in module jax.nn) serial() (in module jax.example_libraries.stax) set_printoptions() (in module jax.numpy) setdiff1d() (in module jax.numpy) setxor1d() (in module jax.numpy) sf() (in module jax.scipy.stats.beta) (in module jax.scipy.stats.cauchy) (in module jax.scipy.stats.chi2) (in module jax.scipy.stats.gamma) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.truncnorm) sgd() (in module jax.example_libraries.optimizers) shape() (in module jax.numpy) shape_dependent() (in module jax.example_libraries.stax) ShapeDtypeStruct (class in jax) shard_shape() (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) (jax.sharding.XLACompatibleSharding method) Sharding (class in jax.sharding) shift_left() (in module jax.lax) shift_right_arithmetic() (in module jax.lax) shift_right_logical() (in module jax.lax) shutdown() (in module jax.distributed) sigmoid() (in module jax.nn) sign() (in module jax.lax) (in module jax.numpy) signbit() (in module jax.numpy) signedinteger (class in jax.numpy) silu() (in module jax.nn) sin() (in module jax.lax) (in module jax.numpy) sinc() (in module jax.numpy) single (in module jax.numpy)

SingleDeviceSharding (class in jax.sharding) sinh() (in module jax.lax) (in module jax.numpy) size() (in module jax.numpy) Slerp (class in jax.scipy.spatial.transform) slice() (in module jax.lax) slice_in_dim() (in module jax.lax) slogdet() (in module jax.numpy.linalg) sm3() (in module jax.example_libraries.optimizers) smallest_normal (jax.numpy.finfo attribute) smallest_subnormal (jax.numpy.finfo attribute) soft_sign() (in module jax.nn) softmax() (in module jax.nn) softplus() (in module jax.nn) solve() (in module jax.numpy.linalg) (in module jax.scipy.linalg) solve_triangular() (in module jax.scipy.linalg) sort() (in module jax.lax) (in module jax.numpy) sort_complex() (in module jax.numpy) sort_key_val() (in module jax.lax) sparsify() (in module jax.experimental.sparse) spence() (in module jax.scipy.special) sph_harm() (in module jax.scipy.special) split() (in module jax.numpy) (in module jax.random) SPMD spsolve() (in module jax.experimental.sparse.linalg) sqrt() (in module jax.lax) (in module jax.numpy) sqrtm() (in module jax.scipy.linalg) square() (in module jax.lax) (in module jax.numpy) squareplus() (in module jax.nn) squeeze() (in module jax.lax) (in module jax.numpy) stack() (in module jax.numpy) standardize() (in module jax.nn) start_server() (in module jax.profiler) start_trace() (in module jax.profiler) static std() (in module jax.numpy) StepTraceAnnotation (class in jax.profiler) stft() (in module jax.scipy.signal) stop_gradient() (in module jax.lax) stop_trace() (in module jax.profiler) StoreException sub() (in module jax.lax) subtract() (in module jax.numpy) subtree_defs (jax.example_libraries.optimizers.OptimizerState attribute) sum() (in module jax.numpy) SumPool() (in module jax.example_libraries.stax) svd() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) swapaxes() (in module jax.numpy) swish() (in module jax.nn) switch() (in module jax.lax) sync_global_devices() (in module jax.experimental.multihost_utils)

T

t() (in module jax.random) take() (in module jax.numpy) take_along_axis() (in module jax.numpy) tan() (in module jax.lax) (in module jax.numpy) tanh() (in module jax.lax) (in module jax.numpy) tensordot() (in module jax.numpy) tensorinv() (in module jax.numpy.linalg) tensorsolve() (in module jax.numpy.linalg) threefry2x32_p (in module jax.extend.random) threefry_2x32() (in module jax.extend.random) threefry_prng_impl (in module jax.extend.random) tie_in() (in module jax.lax) tile() (in module jax.numpy) tiny (jax.numpy.finfo attribute) to_dlpack() (in module jax.dlpack) todense() (in module jax.experimental.sparse) toeplitz() (in module jax.scipy.linalg) top_k() (in module jax.lax) TPU trace() (in module jax.numpy) (in module jax.profiler) TraceAnnotation (class in jax.profiler) Tracer TracerArrayConversionError (class in jax.errors) TracerBoolConversionError (class in jax.errors) TracerIntegerConversionError (class in jax.errors) transfer_guard() (in module jax) transformation (in module jax.extend.linear_util) transformation_with_aux (in module jax.extend.linear_util) transpose() (in module jax.lax) (in module jax.numpy)

trapezoid() (in module jax.scipy.integrate) trapz() (in module jax.numpy) tree_all() (in module jax.tree_util) tree_def (jax.example_libraries.optimizers.OptimizerState attribute) tree_flatten() (in module jax.tree_util) tree_flatten_with_path() (in module jax.tree_util) tree_leaves() (in module jax.tree_util) tree_leaves_with_path() (in module jax.tree_util) tree_map() (in module jax.tree_util) tree_map_with_path() (in module jax.tree_util) tree_reduce() (in module jax.tree_util) tree_structure() (in module jax.tree_util) tree_transpose() (in module jax.tree_util) tree_unflatten() (in module jax.tree_util) treedef_children() (in module jax.tree_util) treedef_is_leaf() (in module jax.tree_util) treedef_tuple() (in module jax.tree_util) tri() (in module jax.numpy) triangular() (in module jax.random) triangular_solve() (in module jax.lax.linalg) tridiagonal() (in module jax.lax.linalg) tridiagonal_solve() (in module jax.lax.linalg) tril() (in module jax.numpy) (in module jax.scipy.linalg) tril_indices() (in module jax.numpy) tril_indices_from() (in module jax.numpy) trim_zeros() (in module jax.numpy) triu() (in module jax.numpy) (in module jax.scipy.linalg) triu_indices() (in module jax.numpy) triu_indices_from() (in module jax.numpy) true_divide() (in module jax.numpy) trunc() (in module jax.numpy) truncated_normal() (in module jax.nn.initializers) (in module jax.random)

U

ufunc (class in jax.numpy) uint (in module jax.numpy) uint16 (class in jax.numpy) uint32 (class in jax.numpy) uint64 (class in jax.numpy) uint8 (class in jax.numpy) UnexpectedTracerError (class in jax.errors) uniform() (in module jax.nn.initializers) (in module jax.random)

union1d() (in module jax.numpy) unique() (in module jax.numpy) unpack_optimizer_state() (in module jax.example_libraries.optimizers) unpackbits() (in module jax.numpy) unravel_index() (in module jax.numpy) unsafe_rbg_prng_impl (in module jax.extend.random) unsignedinteger (class in jax.numpy) unwrap() (in module jax.numpy) update_fn (jax.example_libraries.optimizers.Optimizer attribute) user_checks (in module jax.experimental.checkify)

V

value_and_grad() (in module jax) (in module jax.experimental.sparse) vander() (in module jax.numpy) var() (in module jax.numpy) variance_scaling() (in module jax.nn.initializers) vdot() (in module jax.numpy) vectorize() (in module jax.numpy)

visualize_array_sharding() (in module jax.debug) visualize_sharding() (in module jax.debug) VJP vjp() (in module jax) vmap() (in module jax) vsplit() (in module jax.numpy) vstack() (in module jax.numpy)

W

wald() (in module jax.random) weak type weibull_min() (in module jax.random) welch() (in module jax.scipy.signal) where() (in module jax.numpy) while_loop() (in module jax.lax) with_memory_kind() (jax.sharding.GSPMDSharding method) (jax.sharding.NamedSharding method) (jax.sharding.PmapSharding method) (jax.sharding.PositionalSharding method) (jax.sharding.Sharding method) (jax.sharding.SingleDeviceSharding method)

with_sharding_constraint() (in module jax.lax) wrap_init() (in module jax.extend.linear_util) wrap_key_data() (in module jax.random) Wrapped (class in jax.stages) WrappedFun (class in jax.extend.linear_util)

X

XLA xla_computation() (in module jax) XLACompatibleSharding (class in jax.sharding)

xlog1py (in module jax.scipy.special) xlogy (in module jax.scipy.special) xmap() (in module jax.experimental.maps)

Z

zeros() (in module jax.nn.initializers) (in module jax.numpy) zeros_like() (in module jax.numpy)

zeros_like_array() (in module jax.lax) zeta (in module jax.scipy.special) zeta() (in module jax.lax)