Index

_

__call__() (jax.stages.Compiled method) (jax.stages.Wrapped method) __init__() (jax.Array method) (jax.custom_batching.custom_vmap 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.custom_dce.custom_dce method) (jax.experimental.pallas.BlockSpec method) (jax.experimental.pallas.GridSpec method) (jax.experimental.pallas.MemoryRef method) (jax.experimental.pallas.mosaic_gpu.Barrier method) (jax.experimental.pallas.mosaic_gpu.GPUBlockSpec method) (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams method) (jax.experimental.pallas.mosaic_gpu.GPUMemorySpace method) (jax.experimental.pallas.mosaic_gpu.Layout method) (jax.experimental.pallas.mosaic_gpu.SwizzleTransform method) (jax.experimental.pallas.mosaic_gpu.TilingTransform method) (jax.experimental.pallas.mosaic_gpu.TransposeTransform method) (jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef method) (jax.experimental.pallas.Slice method) (jax.experimental.pallas.triton.TritonCompilerParams 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.export.SymbolicScope method) (jax.extend.core.ClosedJaxpr method) (jax.extend.core.Jaxpr method) (jax.extend.core.JaxprEqn method) (jax.extend.core.Literal method) (jax.extend.core.Primitive method) (jax.extend.core.Token method) (jax.extend.core.Var method) (jax.extend.linear_util.WrappedFun method) (jax.lax.linalg.SvdAlgorithm 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.interpolate.RegularGridInterpolator 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)

_get_vjp (jax.export.Exported attribute)

A

abs() (in module jax.lax) (in module jax.numpy) absolute() (in module jax.numpy) ACC (in module jax.experimental.pallas.mosaic_gpu) accumulation_type (jax.lax.DotAlgorithmPreset property) acos() (in module jax.lax) (in module jax.numpy) acosh() (in module jax.lax) (in module jax.numpy) 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.numpy) add() (in module jax.lax) addressable_devices (jax.sharding.NamedSharding property) (jax.sharding.Sharding property) addressable_devices_indices_map() (jax.sharding.Sharding method) addressable_shards (jax.Array property) after_all() (in module jax.lax) all() (in module jax.numpy) (in module jax.tree) (jax.Array method) 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) (jax.Array method) ANY_F8_ANY_F8_ANY (jax.lax.DotAlgorithmPreset attribute) ANY_F8_ANY_F8_ANY_FAST_ACCUM (jax.lax.DotAlgorithmPreset attribute) ANY_F8_ANY_F8_F32 (jax.lax.DotAlgorithmPreset attribute) ANY_F8_ANY_F8_F32_FAST_ACCUM (jax.lax.DotAlgorithmPreset attribute) append() (in module jax.numpy) apply_along_axis() (in module jax.numpy) apply_over_axes() (in module jax.numpy) approx_math (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) approx_max_k() (in module jax.lax) approx_min_k() (in module jax.lax) approx_tanh() (in module jax.experimental.pallas.triton) 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) (jax.Array method)

argmin() (in module jax.lax) (in module jax.numpy) (jax.Array method) argpartition() (in module jax.numpy) (jax.Array method) argsort() (in module jax.numpy) (jax.Array method) 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) array_types (in module jax.extend.core) 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) (in module jax.numpy) asinh() (in module jax.lax) (in module jax.numpy) assert_equal() (in module jax.experimental.multihost_utils) associative_scan() (in module jax.lax) astype() (in module jax.numpy) (jax.Array method) at (jax.Array property) (jax.numpy.ndarray property) atan() (in module jax.lax) (in module jax.numpy) atan2() (in module jax.lax) (in module jax.numpy) atanh() (in module jax.lax) (in module jax.numpy) atleast_1d() (in module jax.numpy) atleast_2d() (in module jax.numpy) atleast_3d() (in module jax.numpy) atomic_add() (in module jax.experimental.pallas) atomic_and() (in module jax.experimental.pallas) atomic_cas() (in module jax.experimental.pallas) atomic_max() (in module jax.experimental.pallas) atomic_min() (in module jax.experimental.pallas) atomic_or() (in module jax.experimental.pallas) atomic_xchg() (in module jax.experimental.pallas) atomic_xor() (in module jax.experimental.pallas) automatic_checks (in module jax.experimental.checkify) average() (in module jax.numpy) AvgPool() (in module jax.example_libraries.stax) AWAY_FROM_ZERO (jax.lax.RoundingMethod attribute) axis_index() (in module jax.lax) axis_size() (in module jax.lax)

B

ball() (in module jax.random) Barrier (class in jax.experimental.pallas.mosaic_gpu) barrier_arrive() (in module jax.experimental.pallas.mosaic_gpu) barrier_wait() (in module jax.experimental.pallas.mosaic_gpu) 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)

BF16_BF16_BF16 (jax.lax.DotAlgorithmPreset attribute) BF16_BF16_F32 (jax.lax.DotAlgorithmPreset attribute) BF16_BF16_F32_X3 (jax.lax.DotAlgorithmPreset attribute) BF16_BF16_F32_X6 (jax.lax.DotAlgorithmPreset attribute) BF16_BF16_F32_X9 (jax.lax.DotAlgorithmPreset attribute) 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.numpy) bitwise_and() (in module jax.lax) bitwise_count() (in module jax.numpy) bitwise_invert() (in module jax.numpy) bitwise_left_shift() (in module jax.numpy) bitwise_not() (in module jax.lax) (in module jax.numpy) bitwise_or (in module jax.numpy) bitwise_or() (in module jax.lax) bitwise_right_shift() (in module jax.numpy) bitwise_xor (in module jax.numpy) bitwise_xor() (in module jax.lax) blackman() (in module jax.numpy) block() (in module jax.numpy) block_diag() (in module jax.scipy.linalg) block_until_ready() (in module jax) BlockSpec (class in jax.experimental.pallas) bool_ (in module 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.experimental.pallas) (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() (jax.export.Exported method) callback() (in module jax.debug) calling_convention_version (jax.export.Exported attribute) 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.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.norm) (in module jax.scipy.stats.poisson) (in module jax.scipy.stats.truncnorm) (in module jax.scipy.stats.uniform) 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_grads() (in module jax.test_util) check_jvp() (in module jax.test_util) check_tracer_leaks (in module jax) check_vjp() (in module jax.test_util) 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) cholesky_update() (in module jax.lax.linalg) choose() (in module jax.numpy) (jax.Array method) clamp() (in module jax.lax) clear_caches() (in module jax) clip() (in module jax.numpy) (jax.Array method) clip_grads() (in module jax.example_libraries.optimizers) clone() (in module jax.random) ClosedJaxpr (class in jax.extend.core) closure_convert() (in module jax) clz() (in module jax.lax) collapse() (in module jax.lax) column_stack() (in module jax.numpy) commit_smem() (in module jax.experimental.pallas.mosaic_gpu) committed (jax.Array property) 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 composite() (in module jax.lax) compress() (in module jax.numpy) (jax.Array method) concat() (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) (jax.Array method) conjugate() (in module jax.numpy) (jax.Array method) 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) (jax.Array method) copy_gmem_to_smem() (in module jax.experimental.pallas.mosaic_gpu) copy_smem_to_gmem() (in module jax.experimental.pallas.mosaic_gpu) copy_to_host_async() (in module jax) (jax.Array method) 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 create_device_mesh() (in module jax.experimental.mesh_utils) create_hybrid_device_mesh() (in module jax.experimental.mesh_utils) cross() (in module jax.numpy) (in module jax.numpy.linalg) 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) (jax.Array method) cumsum() (in module jax.lax) (in module jax.numpy) (jax.Array method) cumulative_prod() (in module jax.numpy) cumulative_sum() (in module jax.numpy) custom_call() (jax.export.DisabledSafetyCheck class method) custom_dce (class in jax.experimental.custom_dce) 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) custom_vmap (class in jax.custom_batching)

D

data (jax.experimental.sparse.BCOO attribute) dct() (in module jax.scipy.fft) dctn() (in module jax.scipy.fft) debug_barrier() (in module jax.experimental.pallas.triton) debug_infs (in module jax) debug_nans (in module jax) debug_print() (in module jax.experimental.pallas) def_dce() (jax.experimental.custom_dce.custom_dce method) def_vmap() (jax.custom_batching.custom_vmap method) DEFAULT (jax.lax.DotAlgorithmPreset attribute) default() (jax.sharding.PmapSharding class method) default_backend() (in module jax) default_device (in module jax) default_export_platform() (in module jax.export) default_matmul_precision (in module jax) default_prng_impl (in module jax) define_prng_impl() (in module jax.extend.random) defjvp() (jax.custom_jvp method) defjvps() (jax.custom_jvp method) defvjp() (jax.custom_vjp method) deg2rad() (in module jax.numpy) degrees() (in module jax.numpy) delay_release (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) delete() (in module jax.numpy) delta_orthogonal() (in module jax.nn.initializers) Dense() (in module jax.example_libraries.stax) deserialize() (in module jax.export) deserialize_and_load() (in module jax.experimental.serialize_executable) det() (in module jax.numpy.linalg) (in module jax.scipy.linalg) detrend() (in module jax.scipy.signal) Device (class in jax) device (jax.Array property) device_count() (in module jax) device_get() (in module jax) device_memory_profile() (in module jax.profiler) device_put() (in module jax) device_set (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) devices (jax.sharding.PmapSharding property) devices() (in module jax) devices_indices_map() (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) (jax.sharding.SingleDeviceSharding 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) (in module jax.numpy.linalg) (jax.Array method) diff() (in module jax.numpy) digamma() (in module jax.lax) (in module jax.scipy.special) digitize() (in module jax.numpy) dimension_semantics (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) dirichlet() (in module jax.random) disable_jit() (in module jax) disable_x64() (in module jax.experimental) disabled_safety_checks (jax.export.Exported attribute) DisabledSafetyCheck (class in jax.export) 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.experimental.pallas) (in module jax.lax) (in module jax.numpy) (jax.Array method) dot_general() (in module jax.lax) dot_product_attention() (in module jax.nn) DotAlgorithm (class in jax.lax) DotAlgorithmPreset (class in 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.Array property) (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) elementwise_inline_asm() (in module jax.experimental.pallas.triton) elu() (in module jax.nn) emit_pipeline() (in module jax.experimental.pallas.mosaic_gpu) 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.lax) (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) export() (in module jax.export) Exported (class in jax.export) extract() (in module jax.numpy) eye() (in module jax.experimental.sparse) (in module jax.numpy)

F

f() (in module jax.random) F16_F16_F16 (jax.lax.DotAlgorithmPreset attribute) F16_F16_F32 (jax.lax.DotAlgorithmPreset attribute) F32_F32_F32 (jax.lax.DotAlgorithmPreset attribute) F64_F64_F64 (jax.lax.DotAlgorithmPreset attribute) 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) ffi_call() (in module jax.extend.ffi) (in module jax.ffi) ffi_lowering() (in module jax.extend.ffi) (in module jax.ffi) FFT (jax.lax.FftType attribute) 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) FftType (class in jax.lax) fill_diagonal() (in module jax.numpy) finfo (class in jax.numpy) fix() (in module jax.numpy) flat (jax.Array property) flatnonzero() (in module jax.numpy) flatten() (in module jax.tree) (jax.Array method) flatten_with_path() (in module jax.tree) 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) for_tensor_core (jax.experimental.pallas.mosaic_gpu.Barrier attribute) fori_loop() (in module jax.lax) forward-mode autodiff fresnel (in module jax.scipy.special) 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) fun_name (jax.export.Exported attribute) 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) gammasgn() (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) get_scaled_dot_general_config() (in module jax.nn) global_array_to_host_local_array() (in module jax.experimental.multihost_utils) global_shards (jax.Array property) glorot_normal() (in module jax.nn.initializers) glorot_uniform() (in module jax.nn.initializers) glu() (in module jax.nn) GMEM (in module jax.experimental.pallas.mosaic_gpu) gmres() (in module jax.scipy.sparse.linalg) GPU GPUBlockSpec (class in jax.experimental.pallas.mosaic_gpu) GPUCompilerParams (class in jax.experimental.pallas.mosaic_gpu) GPUMemorySpace (class in jax.experimental.pallas.mosaic_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) GridSpec (class in jax.experimental.pallas) 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) has_vjp() (jax.export.Exported method) 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) hilbert() (in module jax.scipy.linalg) 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) (in module jax.scipy.special) i0e() (in module jax.scipy.special) i1() (in module jax.scipy.special) i1e() (in module jax.scipy.special) idct() (in module jax.scipy.fft) idctn() (in module jax.scipy.fft) identity() (in module jax.nn) (in module jax.numpy) iexp (jax.numpy.finfo attribute) IFFT (jax.lax.FftType 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 (jax.Array property) imag() (in module jax.lax) (in module jax.numpy) in_avals (jax.export.Exported attribute) in_shardings_hlo (jax.export.Exported attribute) in_shardings_jax() (jax.export.Exported method) in_tree (jax.export.Exported attribute) (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 (jax.lax.FftType attribute) irfft() (in module jax.numpy.fft) irfft2() (in module jax.numpy.fft) irfftn() (in module jax.numpy.fft) is_custom_call() (jax.export.DisabledSafetyCheck method) is_equivalent_to() (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) is_finite() (in module jax.lax) is_fully_addressable (jax.Array property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) is_fully_replicated (jax.Array property) (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) is_initialized() (in module jax.experimental.compilation_cache.compilation_cache) is_symbolic_dim() (in module jax.export) isclose() (in module jax.numpy) iscomplex() (in module jax.numpy) iscomplexobj() (in module jax.numpy) isdtype() (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) item() (jax.Array method) itemsize (jax.Array property) iterable() (in module jax.numpy) ix_() (in module jax.numpy)

J

jacfwd() (in module jax) jacobian() (in module jax) jacrev() (in module jax) 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.checkify module jax.experimental.compilation_cache.compilation_cache module jax.experimental.custom_dce module jax.experimental.custom_partitioning module jax.experimental.jet module jax.experimental.key_reuse module jax.experimental.mesh_utils module jax.experimental.multihost_utils module jax.experimental.pallas module jax.experimental.pallas.mosaic_gpu module jax.experimental.pallas.tpu module jax.experimental.pallas.triton module jax.experimental.pjit module jax.experimental.serialize_executable module jax.experimental.shard_map module jax.experimental.sparse module jax.experimental.sparse.linalg module jax.export module jax.export.maximum_supported_serialization_version (in module jax.export) jax.export.minimum_supported_serialization_version (in module jax.export) jax.extend module jax.extend.core module jax.extend.core.primitives module jax.extend.ffi 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.ffi 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.cluster.vq module jax.scipy.fft module jax.scipy.integrate module jax.scipy.interpolate 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.test_util module jax.tree module jax.tree_util module jax.typing module jaxpr Jaxpr (class in jax.extend.core) jaxpr_as_fun (in module jax.extend.core) JaxprEqn (class in jax.extend.core) JaxRuntimeError (in module jax.errors) 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) KeyEntry (in module jax.tree_util)

KeyPath (in module jax.tree_util) KeyReuseError (class in jax.errors) 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) Layout (class in jax.experimental.pallas.mosaic_gpu) layout_cast() (in module jax.experimental.pallas.mosaic_gpu) lcm() (in module jax.numpy) ldexp() (in module jax.numpy) le() (in module jax.lax) leaky_relu() (in module jax.nn) leaves() (in module jax.tree) leaves_with_path() (in module jax.tree) 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) Literal (class in jax.extend.core) live_arrays() (in module jax) load() (in module jax.experimental.pallas) (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) (in module jax.scipy.special) 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.expon) (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.expon) (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.Traced method) (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_pivots_to_permutation() (in module jax.lax.linalg) lu_solve() (in module jax.scipy.linalg)

M

machep (jax.numpy.finfo attribute) make_array_from_callback() (in module jax) make_array_from_process_local_data() (in module jax) make_array_from_single_device_arrays() (in module jax) make_jaxpr() (in module jax) make_mesh() (in module jax) make_schedule() (in module jax.example_libraries.optimizers) map() (in module jax.lax) (in module jax.tree) map_coordinates() (in module jax.scipy.ndimage) map_with_path() (in module jax.tree) mask_indices() (in module jax.numpy) matmul() (in module jax.numpy) (in module jax.numpy.linalg) matrix_norm() (in module jax.numpy.linalg) matrix_power() (in module jax.numpy.linalg) matrix_rank() (in module jax.numpy.linalg) matrix_transpose() (in module jax.numpy) (in module jax.numpy.linalg) matvec() (in module jax.numpy) max (jax.numpy.finfo attribute) max() (in module jax.lax) (in module jax.numpy) (jax.Array method) max_concurrent_steps (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) max_contiguous() (in module jax.experimental.pallas) maxexp (jax.numpy.finfo attribute) maximum() (in module jax.numpy) maximum_supported_calling_convention_version (in module jax.export) MaxPool() (in module jax.example_libraries.stax) maxwell() (in module jax.random) mean() (in module jax.numpy) (jax.Array method) median() (in module jax.numpy) memory_analysis() (jax.stages.Compiled method) memory_kind (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) MemoryRef (class in jax.experimental.pallas) merge_linear_aux() (in module jax.extend.linear_util) Mesh (class in jax.sharding) mesh (jax.sharding.NamedSharding property) meshgrid() (in module jax.numpy) mgrid (in module jax.numpy) min (jax.numpy.finfo attribute) min() (in module jax.lax) (in module jax.numpy) (jax.Array method) minexp (jax.numpy.finfo attribute) minimize() (in module jax.scipy.optimize) minimum() (in module jax.numpy) minimum_supported_calling_convention_version (in module jax.export) mish() (in module jax.nn) mlir_module() (jax.export.Exported method) mlir_module_serialized (jax.export.Exported attribute) mod() (in module jax.numpy) mode() (in module jax.scipy.stats) modf() (in module jax.numpy) module jax.debug jax.distributed jax.dlpack jax.dtypes jax.example_libraries jax.example_libraries.optimizers jax.example_libraries.stax jax.experimental.checkify jax.experimental.compilation_cache.compilation_cache jax.experimental.custom_dce jax.experimental.custom_partitioning jax.experimental.jet jax.experimental.key_reuse jax.experimental.mesh_utils jax.experimental.multihost_utils jax.experimental.pallas jax.experimental.pallas.mosaic_gpu jax.experimental.pallas.tpu jax.experimental.pallas.triton jax.experimental.pjit jax.experimental.serialize_executable jax.experimental.shard_map jax.experimental.sparse jax.experimental.sparse.linalg jax.export jax.extend jax.extend.core jax.extend.core.primitives jax.extend.ffi jax.extend.linear_util jax.extend.mlir jax.extend.mlir.dialects jax.extend.mlir.ir jax.extend.mlir.passmanager jax.extend.random jax.ffi 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.cluster.vq jax.scipy.fft jax.scipy.integrate jax.scipy.interpolate 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.test_util jax.tree jax.tree_util jax.typing

module_kept_var_idx (jax.export.Exported attribute) momentum() (in module jax.example_libraries.optimizers) moveaxis() (in module jax.numpy) mT (jax.Array property) mul() (in module jax.lax) multi_dot() (in module jax.numpy.linalg) multigammaln() (in module jax.scipy.special) multinomial() (in module jax.random) multiple_of() (in module jax.experimental.pallas) 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) nbytes (jax.Array property) ndarray (in module jax.numpy) ndim (jax.Array property) 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) no_tracing (in module jax) NonConcreteBooleanIndexError (class in jax.errors) nonzero() (in module jax.numpy) (jax.Array method) norm() (in module jax.numpy.linalg) normal() (in module jax.nn.initializers) (in module jax.random) not_equal() (in module jax.numpy) nr_devices (jax.export.Exported attribute) num_arrivals (jax.experimental.pallas.mosaic_gpu.Barrier attribute) num_barriers (jax.experimental.pallas.mosaic_gpu.Barrier attribute) num_devices (jax.sharding.NamedSharding property) (jax.sharding.PmapSharding property) (jax.sharding.Sharding property) (jax.sharding.SingleDeviceSharding property) num_programs() (in module jax.experimental.pallas) num_stages (jax.experimental.pallas.triton.TritonCompilerParams attribute) num_warps (jax.experimental.pallas.triton.TritonCompilerParams attribute) 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) optimization_barrier() (in module jax.lax) 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) ordered_effects (jax.export.Exported attribute) orthogonal() (in module jax.nn.initializers) (in module jax.random) out_avals (jax.export.Exported attribute) out_shardings_hlo (jax.export.Exported attribute) out_shardings_jax() (jax.export.Exported method) out_tree (jax.export.Exported attribute) outer() (in module jax.numpy) (in module jax.numpy.linalg)

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) pallas_call() (in module jax.experimental.pallas) 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) pascal() (in module jax.scipy.linalg) 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) percentile() (in module jax.numpy) permutation() (in module jax.random) permute_dims() (in module jax.numpy) 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) platform() (jax.export.DisabledSafetyCheck class method) platform_dependent() (in module jax.lax) platforms (jax.export.Exported attribute) 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) positive() (in module jax.numpy) pow() (in module jax.lax) (in module jax.numpy) 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.expon) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) (in module jax.scipy.stats.uniform) Precision (class in jax.lax) precision (jax.numpy.finfo attribute) PrecisionLike (in module jax.lax) primitive Primitive (class in jax.extend.core) 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) process_indices() (in module jax) prod() (in module jax.numpy) (jax.Array method) profile_dir (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) profile_space (jax.experimental.pallas.mosaic_gpu.GPUCompilerParams attribute) program_id() (in module jax.experimental.pallas) 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) (jax.Array method) pure function pure_callback() (in module jax) put() (in module jax.numpy) put_along_axis() (in module jax.numpy) pycapsule() (in module jax.extend.ffi) (in module jax.ffi) pytree

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) RandomAlgorithm (class in jax.lax) rankdata() (in module jax.scipy.stats) ravel() (in module jax.numpy) (jax.Array method) 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 (jax.Array property) real() (in module jax.lax) (in module jax.numpy) reciprocal() (in module jax.lax) (in module jax.numpy) reduce() (in module jax.lax) (in module jax.tree) reduce_and() (in module jax.lax) reduce_max() (in module jax.lax) reduce_min() (in module jax.lax) reduce_or() (in module jax.lax) reduce_precision() (in module jax.lax) reduce_prod() (in module jax.lax) reduce_sum() (in module jax.lax) reduce_window() (in module jax.lax) reduce_xor() (in module jax.lax) register_custom_call_target() (in module jaxlib.xla_client) register_dataclass() (in module jax.tree_util) register_ffi_target() (in module jax.extend.ffi) (in module jax.ffi) register_ffi_type_id() (in module jax.ffi) register_namedtuple_serialization() (in module jax.export) register_pytree_node() (in module jax.tree_util) register_pytree_node_class() (in module jax.tree_util) register_pytree_node_serialization() (in module jax.export) register_pytree_with_keys() (in module jax.tree_util) register_pytree_with_keys_class() (in module jax.tree_util) register_static() (in module jax.tree_util) RegularGridInterpolator (class in jax.scipy.interpolate) 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) (jax.Array method) 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) (jax.Array method) 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 (jax.lax.FftType attribute) 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_DEFAULT (jax.lax.RandomAlgorithm attribute) RNG_PHILOX (jax.lax.RandomAlgorithm attribute) RNG_THREE_FRY (jax.lax.RandomAlgorithm attribute) 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) (jax.Array method) RoundingMethod (class in jax.lax) rsf2csf() (in module jax.scipy.linalg) rsqrt() (in module jax.lax) run_scoped() (in module jax.experimental.pallas) 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) scaled_dot_general() (in module jax.nn) scaled_matmul() (in module jax.nn) 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) (jax.Array method) 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) sem() (in module jax.scipy.stats) sequential_vmap() (in module jax.custom_batching) serial() (in module jax.example_libraries.stax) serialize() (in module jax.experimental.serialize_executable) (jax.export.Exported method) serialized_metadata (jax.experimental.pallas.triton.TritonCompilerParams attribute) set_cache_dir() (in module jax.experimental.compilation_cache.compilation_cache) set_max_registers() (in module jax.experimental.pallas.mosaic_gpu) 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.expon) (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 (jax.Array property) shape() (in module jax.numpy) shape_dependent() (in module jax.example_libraries.stax) ShapeDtypeStruct (class in jax) shard_map() (in module jax) (in module jax.experimental.shard_map) shard_shape() (jax.sharding.PmapSharding method) (jax.sharding.Sharding method) Sharding (class in jax.sharding) sharding (jax.Array property) sharding_spec (jax.sharding.PmapSharding property) 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 (jax.Array property) size() (in module jax.numpy) Slerp (class in jax.scipy.spatial.transform)

Slice (class in jax.experimental.pallas) 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) SMEM (in module jax.experimental.pallas.mosaic_gpu) soft_sign() (in module jax.nn) softmax() (in module jax.nn) (in module jax.scipy.special) 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) (jax.Array method) sort_complex() (in module jax.numpy) sort_key_val() (in module jax.lax) spacing() (in module jax.numpy) sparse_plus() (in module jax.nn) sparse_sigmoid() (in module jax.nn) sparsify() (in module jax.experimental.sparse) spec (jax.sharding.NamedSharding property) spence() (in module jax.scipy.special) sph_harm() (in module jax.scipy.special) split() (in module jax.lax) (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) (jax.Array method) 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) (jax.Array method) StepTraceAnnotation (class in jax.profiler) stft() (in module jax.scipy.signal) stop_gradient() (in module jax.lax) stop_trace() (in module jax.profiler) store() (in module jax.experimental.pallas) StoreException structure() (in module jax.tree) sub() (in module jax.lax) subtract (in module jax.numpy) subtree_defs (jax.example_libraries.optimizers.OptimizerState attribute) sum() (in module jax.numpy) (jax.Array method) SumPool() (in module jax.example_libraries.stax) supported_lhs_types (jax.lax.DotAlgorithmPreset property) supported_output_types() (jax.lax.DotAlgorithmPreset method) supported_rhs_types (jax.lax.DotAlgorithmPreset property) svd() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) SvdAlgorithm (class in jax.lax.linalg) svdvals() (in module jax.numpy.linalg) swap() (in module jax.experimental.pallas) swapaxes() (in module jax.numpy) (jax.Array method) swish() (in module jax.nn) switch() (in module jax.lax) SwizzleTransform (class in jax.experimental.pallas.mosaic_gpu) symbolic_args_specs() (in module jax.export) symbolic_shape() (in module jax.export) SymbolicScope (class in jax.export) symmetric_product() (in module jax.lax.linalg) sync_global_devices() (in module jax.experimental.multihost_utils)

T

T (jax.Array property) t() (in module jax.random) take() (in module jax.numpy) (jax.Array method) 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) (in module jax.numpy.linalg) tensorinv() (in module jax.numpy.linalg) tensorsolve() (in module jax.numpy.linalg) TF32_TF32_F32 (jax.lax.DotAlgorithmPreset attribute) TF32_TF32_F32_X3 (jax.lax.DotAlgorithmPreset attribute) threefry2x32_p (in module jax.extend.random) threefry_2x32() (in module jax.extend.random) threefry_prng_impl (in module jax.extend.random) tile() (in module jax.numpy) TilingTransform (class in jax.experimental.pallas.mosaic_gpu) tiny (jax.numpy.finfo attribute) to_device() (jax.Array method) to_dlpack() (in module jax.dlpack) TO_NEAREST_EVEN (jax.lax.RoundingMethod attribute) todense() (in module jax.experimental.sparse) toeplitz() (in module jax.scipy.linalg) Token (class in jax.extend.core) top_k() (in module jax.lax) TPU trace() (in module jax.numpy) (in module jax.numpy.linalg) (in module jax.profiler) (jax.Array method) (jax.stages.Wrapped method) TraceAnnotation (class in jax.profiler) Traced (class in jax.stages) 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) (in module jax.tree) (jax.Array method) TransposeTransform (class in jax.experimental.pallas.mosaic_gpu) trapezoid() (in module jax.numpy) (in module jax.scipy.integrate) 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) tril_indices() (in module jax.numpy) tril_indices_from() (in module jax.numpy) trim_zeros() (in module jax.numpy) TritonCompilerParams (class in jax.experimental.pallas.triton) triu() (in module jax.numpy) 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) unflatten() (in module jax.tree) uniform() (in module jax.nn.initializers) (in module jax.random) union1d() (in module jax.numpy) unique() (in module jax.numpy) unique_all() (in module jax.numpy)

unique_counts() (in module jax.numpy) unique_inverse() (in module jax.numpy) unique_values() (in module jax.numpy) unordered_effects (jax.export.Exported attribute) 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) unstack() (in module jax.numpy) unwrap() (in module jax.numpy) update_fn (jax.example_libraries.optimizers.Optimizer attribute) user_checks (in module jax.experimental.checkify) uses_global_constants (jax.export.Exported attribute)

V

value_and_grad() (in module jax) (in module jax.experimental.sparse) vander() (in module jax.numpy) Var (class in jax.extend.core) var() (in module jax.numpy) (jax.Array method) variance_scaling() (in module jax.nn.initializers) vdot() (in module jax.numpy) vecdot() (in module jax.numpy) (in module jax.numpy.linalg) vecmat() (in module jax.numpy)

vector_norm() (in module jax.numpy.linalg) vectorize() (in module jax.numpy) view() (jax.Array method) visualize_array_sharding() (in module jax.debug) visualize_sharding() (in module jax.debug) VJP vjp() (in module jax) (jax.export.Exported method) vmap() (in module jax) vq() (in module jax.scipy.cluster.vq) vsplit() (in module jax.numpy) vstack() (in module jax.numpy)

W

wait_smem_to_gmem() (in module jax.experimental.pallas.mosaic_gpu) wald() (in module jax.random) weak type weibull_min() (in module jax.random) welch() (in module jax.scipy.signal) wgmma() (in module jax.experimental.pallas.mosaic_gpu) wgmma_wait() (in module jax.experimental.pallas.mosaic_gpu) WGMMAAccumulatorRef (class in jax.experimental.pallas.mosaic_gpu) when() (in module jax.experimental.pallas) where() (in module jax.numpy)

while_loop() (in module jax.lax) with_memory_kind() (jax.sharding.NamedSharding method) (jax.sharding.PmapSharding 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	xlog1py (in module jax.scipy.special) xlogy (in module jax.scipy.special)

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)