索引

_

__call__() （jax.stages.Compiled 方法） __call__() （jax.stages.Wrapped 方法） __init__() （jax.Array 方法） __init__() （jax.core.ClosedJaxpr 方法） __init__() （jax.core.Jaxpr 方法） __init__() （jax.custom_jvp 方法） __init__() （jax.custom_vjp 方法） __init__() （jax.Device 方法） __init__() （jax.dtypes.bfloat16 方法） __init__() （jax.dtypes.prng_key 方法） __init__() （jax.experimental.checkify.Error 方法） __init__() （jax.experimental.pallas.BlockSpec 方法） __init__() （jax.experimental.pallas.GridSpec 方法） __init__() （jax.experimental.pallas.Slice 方法） __init__() （jax.experimental.sparse.BCOO 方法） __init__() （jax.experimental.sparse.BCSR 方法） __init__() （jax.experimental.sparse.COO 方法） __init__() （jax.experimental.sparse.CSC 方法） __init__() （jax.experimental.sparse.CSR 方法） __init__() （jax.experimental.sparse.JAXSparse 方法） __init__() （jax.export.DisabledSafetyCheck 方法） __init__() （jax.export.Exported 方法） __init__() （jax.export.SymbolicScope 方法） __init__() （jax.extend.linear_util.WrappedFun 方法） __init__() （jax.numpy.character 方法） __init__() （jax.numpy.complex128 方法） __init__() （jax.numpy.complex64 方法） __init__() （jax.numpy.complexfloating 方法） __init__() （jax.numpy.dtype 方法） __init__() （jax.numpy.finfo 方法） __init__() （jax.numpy.flexible 方法）

__init__() （jax.numpy.float16 方法） __init__() （jax.numpy.float32 方法） __init__() （jax.numpy.float64 方法） __init__() （jax.numpy.floating 方法） __init__() （jax.numpy.generic 方法） __init__() （jax.numpy.iinfo 方法） __init__() （jax.numpy.inexact 方法） __init__() （jax.numpy.int16 方法） __init__() （jax.numpy.int32 方法） __init__() （jax.numpy.int64 方法） __init__() （jax.numpy.int8 方法） __init__() （jax.numpy.integer 方法） __init__() （jax.numpy.number 方法） __init__() （jax.numpy.object_ 方法） __init__() （jax.numpy.signedinteger 方法） __init__() （jax.numpy.ufunc 方法） __init__() （jax.numpy.uint16 方法） __init__() （jax.numpy.uint32 方法） __init__() （jax.numpy.uint64 方法） __init__() （jax.numpy.uint8 方法） __init__() （jax.numpy.unsignedinteger 方法） __init__() （jax.profiler.StepTraceAnnotation 方法） __init__() （jax.profiler.TraceAnnotation 方法） __init__() （jax.scipy.interpolate.RegularGridInterpolator 方法） __init__() （jax.scipy.optimize.OptimizeResults 方法） __init__() （jax.scipy.spatial.transform.Rotation 方法） __init__() （jax.scipy.spatial.transform.Slerp 方法） __init__() （jax.scipy.stats.gaussian_kde 方法） __init__() （jax.ShapeDtypeStruct 方法） __init__() （jax.tree_util.Partial 方法） _get_vjp（jax.export.Exported 属性）

A

abs()（在 jax.lax 模块中） abs()（在 jax.numpy 模块中） absolute()（在 jax.numpy 模块中） acos()（在 jax.lax 模块中） acos()（在 jax.numpy 模块中） acosh()（在 jax.lax 模块中） acosh()（在 jax.numpy 模块中） adagrad()（在 jax.example_libraries.optimizers 模块中） adam()（在 jax.example_libraries.optimizers 模块中） adamax()（在 jax.example_libraries.optimizers 模块中） add()（在 jax.lax 模块中） add()（在 jax.numpy 模块中） addressable_devices_indices_map() （jax.sharding.Sharding 方法） addressable_devices（jax.sharding.NamedSharding 属性） addressable_devices（jax.sharding.Sharding 属性） addressable_shards（jax.Array 属性） after_all()（在 jax.lax 模块中） all() （jax.Array 方法） all()（在 jax.numpy 模块中） all()（在 jax.tree 模块中） all_checks()（在 jax.experimental.checkify 模块中） all_gather()（在 jax.lax 模块中） all_leaves()（在 jax.tree_util 模块中） all_to_all()（在 jax.lax 模块中） allclose()（在 jax.numpy 模块中） amax()（在 jax.numpy 模块中） amin()（在 jax.numpy 模块中） angle()（在 jax.numpy 模块中） annotate_function()（在 jax.profiler 模块中） any() （jax.Array 方法） any()（在 jax.numpy 模块中） append()（在 jax.numpy 模块中） apply_along_axis()（在 jax.numpy 模块中） apply_over_axes()（在 jax.numpy 模块中） approx_max_k()（在 jax.lax 模块中） approx_min_k()（在 jax.lax 模块中） arange()（在 jax.numpy 模块中） arccos()（在 jax.numpy 模块中） arccosh()（在 jax.numpy 模块中） arcsin()（在 jax.numpy 模块中） arcsinh()（在 jax.numpy 模块中） arctan()（在 jax.numpy 模块中） arctan2()（在 jax.numpy 模块中） arctanh()（在 jax.numpy 模块中） argmax() （jax.Array 方法） argmax()（在 jax.lax 模块中） argmax()（在 jax.numpy 模块中） argmin() （jax.Array 方法）

argmin()（在 jax.lax 模块中） argmin()（在 jax.numpy 模块中） argpartition() （jax.Array 方法） argpartition()（在 jax.numpy 模块中） argsort() （jax.Array 方法） argsort()（在 jax.numpy 模块中） argwhere()（在 jax.numpy 模块中） around()（在 jax.numpy 模块中） array()（在 jax.numpy 模块中） array_equal()（在 jax.numpy 模块中） array_equiv()（在 jax.numpy 模块中） array_repr()（在 jax.numpy 模块中） array_split()（在 jax.numpy 模块中） array_str()（在 jax.numpy 模块中） ArrayLike()（在 jax.typing 模块中） Array（jax 中的类） as_text() （jax.stages.Compiled 方法） as_text() （jax.stages.Lowered 方法） asarray()（在 jax.numpy 模块中） asin()（在 jax.lax 模块中） asin()（在 jax.numpy 模块中） asinh()（在 jax.lax 模块中） asinh()（在 jax.numpy 模块中） assert_equal()（在 jax.experimental.multihost_utils 模块中） associative_scan()（在 jax.lax 模块中） astype() （jax.Array 方法） astype()（在 jax.numpy 模块中） atan()（在 jax.lax 模块中） atan()（在 jax.numpy 模块中） atan2()（在 jax.lax 模块中） atan2()（在 jax.numpy 模块中） atanh()（在 jax.lax 模块中） atanh()（在 jax.numpy 模块中） atleast_1d()（在 jax.numpy 模块中） atleast_2d()（在 jax.numpy 模块中） atleast_3d()（在 jax.numpy 模块中） atomic_add()（在 jax.experimental.pallas 模块中） atomic_and()（在 jax.experimental.pallas 模块中） atomic_cas()（在 jax.experimental.pallas 模块中） atomic_max()（在 jax.experimental.pallas 模块中） atomic_min()（在 jax.experimental.pallas 模块中） atomic_or()（在 jax.experimental.pallas 模块中） atomic_xchg()（在 jax.experimental.pallas 模块中） at（jax.Array 属性） at（jax.numpy.ndarray 属性） automatic_checks()（在 jax.experimental.checkify 模块中） average()（在 jax.numpy 模块中） AvgPool()（在 jax.example_libraries.stax 模块中） axis_index()（在 jax.lax 模块中）

B

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

betainc()（在 jax.scipy.special 模块中） betaln()（在 jax.scipy.special 模块中） bfloat16（jax.dtypes 中的类） bicgstab()（在 jax.scipy.sparse.linalg 模块中） bincount()（在 jax.numpy 模块中） binomial()（在 jax.random 模块中） bitcast_convert_type()（在 jax.lax 模块中） bits()（在 jax.random 模块中） bits（jax.numpy.finfo 属性） bitwise_and()（在 jax.lax 模块中） bitwise_and()（在 jax.numpy 模块中） bitwise_count()（在 jax.numpy 模块中） bitwise_invert()（在 jax.numpy 模块中） bitwise_left_shift()（在 jax.numpy 模块中） bitwise_not()（在 jax.lax 模块中） bitwise_not()（在 jax.numpy 模块中） bitwise_or()（在 jax.lax 模块中） bitwise_or()（在 jax.numpy 模块中） bitwise_right_shift()（在 jax.numpy 模块中） bitwise_xor()（在 jax.lax 模块中） bitwise_xor()（在 jax.numpy 模块中） blackman()（在 jax.numpy 模块中） block()（在 jax.numpy 模块中） block_diag()（在 jax.scipy.linalg 模块中） block_until_ready()（在 jax 模块中） BlockSpec（jax.experimental.pallas 中的类） bool_()（在 jax.numpy 模块中） breakpoint()（在 jax.debug 模块中） broadcast()（在 jax.lax 模块中） broadcast_arrays()（在 jax.numpy 模块中） broadcast_in_dim()（在 jax.lax 模块中） broadcast_one_to_all()（在 jax.experimental.multihost_utils 模块中） broadcast_shapes()（在 jax.lax 模块中） broadcast_shapes()（在 jax.numpy 模块中） broadcast_to()（在 jax.numpy 模块中） broadcast_to_rank()（在 jax.lax 模块中） broadcasted_iota()（在 jax.lax 模块中） build_tree()（在 jax.tree_util 模块中）

C

c_()（在 jax.numpy 模块中） cache()（在 jax.extend.linear_util 模块中） call()（在 jax.experimental.host_callback 模块中） callback()（在 jax.debug 模块中） CallbackException calling_convention_version（jax.export.Exported 属性） can_cast()（在 jax.numpy 模块中） canonicalize_dtype()（在 jax.dtypes 模块中） categorical()（在 jax.random 模块中） cauchy()（在 jax.random 模块中） cbrt()（在 jax.lax 模块中） cbrt()（在 jax.numpy 模块中） cdf()（在 jax.scipy.stats.bernoulli 模块中） cdf()（在 jax.scipy.stats.beta 模块中） cdf()（在 jax.scipy.stats.cauchy 模块中） cdf()（在 jax.scipy.stats.chi2 模块中） cdf()（在 jax.scipy.stats.gamma 模块中） cdf()（在 jax.scipy.stats.gennorm 模块中） cdf()（在 jax.scipy.stats.laplace 模块中） cdf()（在 jax.scipy.stats.logistic 模块中） cdf()（在 jax.scipy.stats.norm 模块中） cdf()（在 jax.scipy.stats.poisson 模块中） cdf()（在 jax.scipy.stats.truncnorm 模块中） cdf()（在 jax.scipy.stats.uniform 模块中） cdouble()（在 jax.numpy 模块中） ceil()（在 jax.lax 模块中） ceil()（在 jax.numpy 模块中） celu()（在 jax.nn 模块中） cg()（在 jax.scipy.sparse.linalg 模块中） character（jax.numpy 中的类） check()（在 jax.experimental.checkify 模块中） check_error()（在 jax.experimental.checkify 模块中） check_tracer_leaks()（在 jax 模块中） checkify()（在 jax.experimental.checkify 模块中） checking_leaks()（在 jax 模块中） checkpoint()（在 jax 模块中） chisquare()（在 jax.random 模块中） cho_factor()（在 jax.scipy.linalg 模块中） cho_solve()（在 jax.scipy.linalg 模块中） choice()（在 jax.random 模块中） cholesky()（在 jax.lax.linalg 模块中） cholesky()（在 jax.numpy.linalg 模块中） cholesky()（在 jax.scipy.linalg 模块中） choose() （jax.Array 方法） choose()（在 jax.numpy 模块中） clamp()（在 jax.lax 模块中） clear_caches()（在 jax 模块中） clip() （jax.Array 方法） clip()（在 jax.numpy 模块中） clip_grads()（在 jax.example_libraries.optimizers 模块中） clone()（在 jax.random 模块中） ClosedJaxpr（jax.core 中的类） closure_convert()（在 jax 模块中） clz()（在 jax.lax 模块中） collapse()（在 jax.lax 模块中） column_stack()（在 jax.numpy 模块中） compile() （jax.stages.Lowered 方法） Compiled（jax.stages 中的类） compiler_ir() （jax.stages.Lowered 方法） complex()（在 jax.lax 模块中） complex128（jax.numpy 中的类） complex64（jax.numpy 中的类） complex_()（在 jax.numpy 模块中） complexfloating（jax.numpy 中的类） ComplexWarning compress() （jax.Array 方法） compress()（在 jax.numpy 模块中） concat()（在 jax.numpy 模块中） concatenate()（在 jax.lax 模块中） concatenate()（在 jax.numpy 模块中） ConcretizationTypeError（jax.errors 中的类） cond()（在 jax.lax 模块中） cond()（在 jax.numpy.linalg 模块中）

config()（在 jax 模块中） conj() （jax.Array 方法） conj()（在 jax.lax 模块中） conj()（在 jax.numpy 模块中） conjugate() （jax.Array 方法） conjugate()（在 jax.numpy 模块中） constant()（在 jax.example_libraries.optimizers 模块中） constant()（在 jax.nn.initializers 模块中） Conv()（在 jax.example_libraries.stax 模块中） conv()（在 jax.lax 模块中） Conv1DTranspose()（在 jax.example_libraries.stax 模块中） conv_dimension_numbers()（在 jax.lax 模块中） conv_general_dilated()（在 jax.lax 模块中） conv_general_dilated_local()（在 jax.lax 模块中） conv_general_dilated_patches()（在 jax.lax 模块中） conv_transpose()（在 jax.lax 模块中） conv_with_general_padding()（在 jax.lax 模块中） ConvDimensionNumbers（jax.lax 中的类） convert_element_type()（在 jax.lax 模块中） ConvGeneralDilatedDimensionNumbers()（在 jax.lax 模块中） convolve()（在 jax.numpy 模块中） convolve()（在 jax.scipy.signal 模块中） convolve2d()（在 jax.scipy.signal 模块中） ConvTranspose()（在 jax.example_libraries.stax 模块中） coo_fromdense()（在 jax.experimental.sparse 模块中） coo_matmat()（在 jax.experimental.sparse 模块中） coo_matvec()（在 jax.experimental.sparse 模块中） coo_todense()（在 jax.experimental.sparse 模块中） COO（jax.experimental.sparse 中的类） copy() （jax.Array 方法） copy()（在 jax.numpy 模块中） copy_to_host_async() （jax.Array 方法） copysign()（在 jax.numpy 模块中） corrcoef()（在 jax.numpy 模块中） correlate()（在 jax.numpy 模块中） correlate()（在 jax.scipy.signal 模块中） correlate2d()（在 jax.scipy.signal 模块中） cos()（在 jax.lax 模块中） cos()（在 jax.numpy 模块中） cosh()（在 jax.lax 模块中） cosh()（在 jax.numpy 模块中） cost_analysis() （jax.stages.Compiled 方法） cost_analysis() （jax.stages.Lowered 方法） count_nonzero()（在 jax.numpy 模块中） cov()（在 jax.numpy 模块中） CPU create_device_mesh()（在 jax.experimental.mesh_utils 模块中） create_hybrid_device_mesh()（在 jax.experimental.mesh_utils 模块中） cross()（在 jax.numpy 模块中） cross()（在 jax.numpy.linalg 模块中） CSC（jax.experimental.sparse 中的类） csd()（在 jax.scipy.signal 模块中） csingle()（在 jax.numpy 模块中） csr_fromdense()（在 jax.experimental.sparse 模块中） csr_matmat()（在 jax.experimental.sparse 模块中） csr_matvec()（在 jax.experimental.sparse 模块中） csr_todense()（在 jax.experimental.sparse 模块中） CSR（jax.experimental.sparse 中的类） cumlogsumexp()（在 jax.lax 模块中） cummax()（在 jax.lax 模块中） cummin()（在 jax.lax 模块中） cumprod() （jax.Array 方法） cumprod()（在 jax.lax 模块中） cumprod()（在 jax.numpy 模块中） cumsum() （jax.Array 方法） cumsum()（在 jax.lax 模块中） cumsum()（在 jax.numpy 模块中） cumulative_sum()（在 jax.numpy 模块中） custom_gradient()（在 jax 模块中） custom_jvp（jax 中的类） custom_linear_solve()（在 jax.lax 模块中） custom_partitioning()（在 jax.experimental.custom_partitioning 模块中） custom_root()（在 jax.lax 模块中） custom_vjp（jax 中的类）

D

data（jax.experimental.sparse.BCOO 属性） dct()（在 jax.scipy.fft 模块中） dctn()（在 jax.scipy.fft 模块中） debug_infs()（在 jax 模块中） debug_nans()（在 jax 模块中） debug_print()（在 jax.experimental.pallas 模块中） default()（jax.sharding.PmapSharding 类方法） default_backend()（在 jax 模块中） default_backend()（在 jax.lib.xla_bridge 模块中） default_device()（在 jax 模块中） default_export_platform()（在 jax.export 模块中） default_matmul_precision()（在 jax 模块中） default_prng_impl()（在 jax 模块中） define_prng_impl()（在 jax.extend.random 模块中） defjvp() （jax.custom_jvp 方法） defjvps() （jax.custom_jvp 方法） defvjp() （jax.custom_vjp 方法） deg2rad()（在 jax.numpy 模块中） degrees()（在 jax.numpy 模块中） delete()（在 jax.numpy 模块中） delta_orthogonal()（在 jax.nn.initializers 模块中） Dense()（在 jax.example_libraries.stax 模块中） deserialize()（在 jax.export 模块中） deserialize_and_load()（在 jax.experimental.serialize_executable 模块中） det()（在 jax.numpy.linalg 模块中） det()（在 jax.scipy.linalg 模块中） detrend()（在 jax.scipy.signal 模块中） device_count()（在 jax 模块中） device_get()（在 jax 模块中） device_memory_profile()（在 jax.profiler 模块中） device_put()（在 jax 模块中） device_put_replicated()（在 jax 模块中） device_put_sharded()（在 jax 模块中） device_set（jax.sharding.GSPMDSharding 属性） device_set（jax.sharding.NamedSharding 属性） device_set（jax.sharding.PmapSharding 属性） device_set（jax.sharding.PositionalSharding 属性） device_set（jax.sharding.Sharding 属性） device_set（jax.sharding.SingleDeviceSharding 属性） devices()（在 jax 模块中） devices_indices_map() （jax.sharding.PmapSharding 方法） devices_indices_map() （jax.sharding.Sharding 方法） devices_indices_map() （jax.sharding.SingleDeviceSharding 方法）

devices（jax.sharding.PmapSharding 属性） Device（jax 中的类） device（jax.Array 属性） diag()（在 jax.numpy 模块中） diag_indices()（在 jax.numpy 模块中） diag_indices_from()（在 jax.numpy 模块中） diagflat()（在 jax.numpy 模块中） diagonal() （jax.Array 方法） diagonal()（在 jax.numpy 模块中） diagonal()（在 jax.numpy.linalg 模块中） diff()（在 jax.numpy 模块中） digamma()（在 jax.lax 模块中） digamma()（在 jax.scipy.special 模块中） digitize()（在 jax.numpy 模块中） dirichlet()（在 jax.random 模块中） disable_jit()（在 jax 模块中） disable_x64()（在 jax.experimental 模块中） disabled_safety_checks（jax.export.Exported 属性） DisabledSafetyCheck（jax.export 中的类） div()（在 jax.lax 模块中） div_checks()（在 jax.experimental.checkify 模块中） divide()（在 jax.numpy 模块中） divmod()（在 jax.numpy 模块中） dot() （jax.Array 方法） dot()（在 jax.lax 模块中） dot()（在 jax.numpy 模块中） dot_general()（在 jax.lax 模块中） dot_product_attention()（在 jax.nn 模块中） double()（在 jax.numpy 模块中） double_sided_maxwell()（在 jax.random 模块中） Dropout()（在 jax.example_libraries.stax 模块中） dsplit()（在 jax.numpy 模块中） dstack()（在 jax.numpy 模块中） DTypeLike()（在 jax.typing 模块中） dtype（jax.Array 属性） dtype（jax.numpy 中的类） dtype（jax.numpy.finfo 属性） dynamic_index_in_dim()（在 jax.lax 模块中） dynamic_slice()（在 jax.lax 模块中） dynamic_slice_in_dim()（在 jax.lax 模块中） dynamic_update_index_in_dim()（在 jax.lax 模块中） dynamic_update_slice()（在 jax.lax 模块中） dynamic_update_slice_in_dim()（在 jax.lax 模块中）

E

ediff1d()（在 jax.numpy 模块中） eig()（在 jax.lax.linalg 模块中） eig()（在 jax.numpy.linalg 模块中） eigh()（在 jax.lax.linalg 模块中） eigh()（在 jax.numpy.linalg 模块中） eigh()（在 jax.scipy.linalg 模块中） eigh_tridiagonal()（在 jax.scipy.linalg 模块中） eigvals()（在 jax.numpy.linalg 模块中） eigvalsh()（在 jax.numpy.linalg 模块中） einsum()（在 jax.numpy 模块中） einsum_path()（在 jax.numpy 模块中） elementwise()（在 jax.example_libraries.stax 模块中） elu()（在 jax.nn 模块中） empty()（在 jax.experimental.sparse 模块中） empty()（在 jax.numpy 模块中） empty_like()（在 jax.numpy 模块中） enable_checks()（在 jax 模块中） enable_custom_prng()（在 jax 模块中） enable_custom_vjp_by_custom_transpose()（在 jax 模块中） enable_x64()（在 jax.experimental 模块中） ensure_compile_time_eval()（在 jax 模块中） entr()（在 jax.scipy.special 模块中） epsneg（jax.numpy.finfo 属性） eps（jax.numpy.finfo 属性） eq()（在 jax.lax 模块中） equal()（在 jax.numpy 模块中） erf()（在 jax.lax 模块中）

erf()（在 jax.scipy.special 模块中） erf_inv()（在 jax.lax 模块中） erfc()（在 jax.lax 模块中） erfc()（在 jax.scipy.special 模块中） erfinv()（在 jax.scipy.special 模块中） Error（jax.experimental.checkify 中的类） eval_shape()（在 jax 模块中） evaluate() （jax.scipy.stats.gaussian_kde 方法） exp()（在 jax.lax 模块中） exp()（在 jax.numpy 模块中） exp1()（在 jax.scipy.special 模块中） exp2()（在 jax.numpy 模块中） expand_dims()（在 jax.lax 模块中） expand_dims()（在 jax.numpy 模块中） expi()（在 jax.scipy.special 模块中） expit()（在 jax.scipy.special 模块中） expm()（在 jax.scipy.linalg 模块中） expm1()（在 jax.lax 模块中） expm1()（在 jax.numpy 模块中） expm_frechet()（在 jax.scipy.linalg 模块中） expn()（在 jax.scipy.special 模块中） exponential()（在 jax.random 模块中） exponential_decay()（在 jax.example_libraries.optimizers 模块中） export()（在 jax.export 模块中） Exported（jax.export 中的类） extract()（在 jax.numpy 模块中） eye()（在 jax.experimental.sparse 模块中） eye()（在 jax.numpy 模块中）

F

f()（在 jax.random 模块中） fabs()（在 jax.numpy 模块中） factorial()（在 jax.scipy.special 模块中） FanInConcat()（在 jax.example_libraries.stax 模块中） FanOut()（在 jax.example_libraries.stax 模块中） ffi_call()（在 jax.extend.ffi 模块中） ffi_lowering()（在 jax.extend.ffi 模块中） fft()（在 jax.lax 模块中） fft()（在 jax.numpy.fft 模块中） fft2()（在 jax.numpy.fft 模块中） fftconvolve()（在 jax.scipy.signal 模块中） fftfreq()（在 jax.numpy.fft 模块中） fftn()（在 jax.numpy.fft 模块中） fftshift()（在 jax.numpy.fft 模块中） fill_diagonal()（在 jax.numpy 模块中） finfo（jax.numpy 中的类） fix()（在 jax.numpy 模块中） flatnonzero()（在 jax.numpy 模块中） flatten() （jax.Array 方法） flatten()（在 jax.tree 模块中） flat（jax.Array 属性） flexible（jax.numpy 中的类） flip()（在 jax.numpy 模块中） fliplr()（在 jax.numpy 模块中） flipud()（在 jax.numpy 模块中） float0()（在 jax.dtypes 模块中） float16（jax.numpy 中的类） float32（jax.numpy 中的类）

float64（jax.numpy 中的类） float_()（在 jax.numpy 模块中） float_checks()（在 jax.experimental.checkify 模块中） float_power()（在 jax.numpy 模块中） floating（jax.numpy 中的类） floor()（在 jax.lax 模块中） floor()（在 jax.numpy 模块中） floor_divide()（在 jax.numpy 模块中） fmax()（在 jax.numpy 模块中） fmin()（在 jax.numpy 模块中） fmod()（在 jax.numpy 模块中） fold_in()（在 jax.random 模块中） fori_loop()（在 jax.lax 模块中） fresnel()（在 jax.scipy.special 模块中） frexp()（在 jax.numpy 模块中） from_dlpack()（在 jax.dlpack 模块中） from_dlpack()（在 jax.numpy 模块中） frombuffer()（在 jax.numpy 模块中） fromfile()（在 jax.numpy 模块中） fromfunction()（在 jax.numpy 模块中） fromiter()（在 jax.numpy 模块中） frompyfunc()（在 jax.numpy 模块中） fromstring()（在 jax.numpy 模块中） full()（在 jax.lax 模块中） full()（在 jax.numpy 模块中） full_like()（在 jax.lax 模块中） full_like()（在 jax.numpy 模块中） fun_name（jax.export.Exported 属性） funm()（在 jax.scipy.linalg 模块中）

G

gamma()（在 jax.random 模块中） gamma()（在 jax.scipy.special 模块中） gammainc()（在 jax.scipy.special 模块中） gammaincc()（在 jax.scipy.special 模块中） gammaln()（在 jax.scipy.special 模块中） gammasgn()（在 jax.scipy.special 模块中） gather()（在 jax.lax 模块中） GatherDimensionNumbers（jax.lax 中的类） GatherScatterMode（jax.lax 中的类） gaussian_kde（jax.scipy.stats 中的类） gcd()（在 jax.numpy 模块中） ge()（在 jax.lax 模块中） gelu()（在 jax.nn 模块中） GeneralConv()（在 jax.example_libraries.stax 模块中） GeneralConvTranspose()（在 jax.example_libraries.stax 模块中） generalized_normal()（在 jax.random 模块中） generic（jax.numpy 中的类） geometric()（在 jax.random 模块中） geomspace()（在 jax.numpy 模块中）

get_backend()（在 jax.lib.xla_bridge 模块中） get_compile_options()（在 jax.lib.xla_bridge 模块中） get_printoptions()（在 jax.numpy 模块中） global_array_to_host_local_array()（在 jax.experimental.multihost_utils 模块中） global_shards（jax.Array 属性） glorot_normal()（在 jax.nn.initializers 模块中） glorot_uniform()（在 jax.nn.initializers 模块中） glu()（在 jax.nn 模块中） gmres()（在 jax.scipy.sparse.linalg 模块中） GPU grad()（在 jax 模块中） grad()（在 jax.experimental.sparse 模块中） gradient()（在 jax.numpy 模块中） greater()（在 jax.numpy 模块中） greater_equal()（在 jax.numpy 模块中） GridSpec（jax.experimental.pallas 中的类） GSPMDSharding（jax.sharding 中的类） gt()（在 jax.lax 模块中） gumbel()（在 jax.random 模块中）

H

hamming()（在 jax.numpy 模块中） hanning()（在 jax.numpy 模块中） hard_sigmoid()（在 jax.nn 模块中） hard_silu()（在 jax.nn 模块中） hard_swish()（在 jax.nn 模块中） hard_tanh()（在 jax.nn 模块中） he_normal()（在 jax.nn.initializers 模块中） he_uniform()（在 jax.nn.initializers 模块中） heaviside()（在 jax.numpy 模块中） hessenberg()（在 jax.lax.linalg 模块中） hessenberg()（在 jax.scipy.linalg 模块中） hessian()（在 jax 模块中）

hfft()（在 jax.numpy.fft 模块中） hilbert()（在 jax.scipy.linalg 模块中） histogram()（在 jax.numpy 模块中） histogram2d()（在 jax.numpy 模块中） histogram_bin_edges()（在 jax.numpy 模块中） histogramdd()（在 jax.numpy 模块中） host_local_array_to_global_array()（在 jax.experimental.multihost_utils 模块中） householder_product()（在 jax.lax.linalg 模块中） hsplit()（在 jax.numpy 模块中） hstack()（在 jax.numpy 模块中） hyp1f1()（在 jax.scipy.special 模块中） hypot()（在 jax.numpy 模块中）

I

i0()（在 jax.numpy 模块中） i0()（在 jax.scipy.special 模块中） i0e()（在 jax.scipy.special 模块中） i1()（在 jax.scipy.special 模块中） i1e()（在 jax.scipy.special 模块中） id_print()（在 jax.experimental.host_callback 模块中） id_tap()（在 jax.experimental.host_callback 模块中） idct()（在 jax.scipy.fft 模块中） idctn()（在 jax.scipy.fft 模块中） identity()（在 jax.numpy 模块中） iexp（jax.numpy.finfo 属性） ifft()（在 jax.numpy.fft 模块中） ifft2()（在 jax.numpy.fft 模块中） ifftn()（在 jax.numpy.fft 模块中） ifftshift()（在 jax.numpy.fft 模块中） igamma()（在 jax.lax 模块中） igammac()（在 jax.lax 模块中） ihfft()（在 jax.numpy.fft 模块中） iinfo（jax.numpy 中的类） imag()（在 jax.lax 模块中） imag()（在 jax.numpy 模块中） imag（jax.Array 属性） in_avals（jax.export.Exported 属性） in_shardings_hlo（jax.export.Exported 属性） in_tree（jax.export.Exported 属性） in_tree（jax.stages.Compiled 属性） in_tree（jax.stages.Lowered 属性） index_checks()（在 jax.experimental.checkify 模块中） index_exp()（在 jax.numpy 模块中） index_in_dim()（在 jax.lax 模块中） index_take()（在 jax.lax 模块中） indices()（在 jax.numpy 模块中） indices（jax.experimental.sparse.BCOO 属性） inexact（jax.numpy 中的类） init_fn（jax.example_libraries.optimizers.Optimizer 属性） initialize()（在 jax.distributed 模块中） initialize_cache()（在 jax.experimental.compilation_cache.compilation_cache 模块中） inner()（在 jax.numpy 模块中） insert()（在 jax.numpy 模块中） inspect_array_sharding()（在 jax.debug 模块中） int16（jax.numpy 中的类） int32（jax.numpy 中的类） int64（jax.numpy 中的类） int8（jax.numpy 中的类） int_()（在 jax.numpy 模块中） integer_pow()（在 jax.lax 模块中） integer（jax.numpy 中的类） integrate_box_1d() （jax.scipy.stats.gaussian_kde 方法） integrate_gaussian() （jax.scipy.stats.gaussian_kde 方法） integrate_kde() （jax.scipy.stats.gaussian_kde 方法） interp()（在 jax.numpy 模块中）

intersect1d()（在 jax.numpy 模块中） inv()（在 jax.numpy.linalg 模块中） inv()（在 jax.scipy.linalg 模块中） inverse_time_decay()（在 jax.example_libraries.optimizers 模块中） invert()（在 jax.numpy 模块中） io_callback()（在 jax.experimental 模块中） iota()（在 jax.lax 模块中） irfft()（在 jax.numpy.fft 模块中） irfft2()（在 jax.numpy.fft 模块中） irfftn()（在 jax.numpy.fft 模块中） is_equivalent_to() （jax.sharding.PmapSharding 方法） is_equivalent_to() （jax.sharding.Sharding 方法） is_finite()（在 jax.lax 模块中） is_fully_addressable（jax.Array 属性） is_fully_addressable（jax.sharding.GSPMDSharding 属性） is_fully_addressable（jax.sharding.NamedSharding 属性） is_fully_addressable（jax.sharding.PmapSharding 属性） is_fully_addressable（jax.sharding.PositionalSharding 属性） is_fully_addressable（jax.sharding.Sharding 属性） is_fully_addressable（jax.sharding.SingleDeviceSharding 属性） is_fully_replicated（jax.Array 属性） is_fully_replicated（jax.sharding.GSPMDSharding 属性） is_fully_replicated（jax.sharding.NamedSharding 属性） is_fully_replicated（jax.sharding.PmapSharding 属性） is_fully_replicated（jax.sharding.PositionalSharding 属性） is_fully_replicated（jax.sharding.Sharding 属性） is_fully_replicated（jax.sharding.SingleDeviceSharding 属性） is_initialized()（在 jax.experimental.compilation_cache.compilation_cache 模块中） is_symbolic_dim()（在 jax.export 模块中） isclose()（在 jax.numpy 模块中） iscomplex()（在 jax.numpy 模块中） iscomplexobj()（在 jax.numpy 模块中） isdtype()（在 jax.numpy 模块中） isf()（在 jax.scipy.stats.cauchy 模块中） isf()（在 jax.scipy.stats.logistic 模块中） isf()（在 jax.scipy.stats.norm 模块中） isfinite()（在 jax.numpy 模块中） isin()（在 jax.numpy 模块中） isinf()（在 jax.numpy 模块中） isnan()（在 jax.numpy 模块中） isneginf()（在 jax.numpy 模块中） isposinf()（在 jax.numpy 模块中） isreal()（在 jax.numpy 模块中） isrealobj()（在 jax.numpy 模块中） isscalar()（在 jax.numpy 模块中） issubdtype()（在 jax.dtypes 模块中） issubdtype()（在 jax.numpy 模块中） istft()（在 jax.scipy.signal 模块中） item() （jax.Array 方法） itemsize（jax.Array 属性） iterable()（在 jax.numpy 模块中） ix_()（在 jax.numpy 模块中）

J

jacfwd()（在 jax 模块中） jacrev()（在 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.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.key_reuse module jax.experimental.mesh_utils module jax.experimental.multihost_utils module jax.experimental.pallas 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()（在 jax.export 模块中） jax.export.minimum_supported_serialization_version()（在 jax.export 模块中） jax.extend 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.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.tree module jax.tree_util module jax.typing module jaxpr Jaxpr（jax.core 中的类） JaxRuntimeError JAXSparse（jax.experimental.sparse 中的类） jet()（在 jax.experimental.jet 模块中） JIT jit()（在 jax 模块中） JoinPoint（jax.example_libraries.optimizers 中的类） JVP jvp()（在 jax 模块中）

K

kaiser()（在 jax.numpy 模块中） key()（在 jax.random 模块中） key_data()（在 jax.random 模块中） KeyEntry()（在 jax.tree_util 模块中）

KeyPath()（在 jax.tree_util 模块中） KeyReuseError（jax.errors 中的类） keystr()（在 jax.tree_util 模块中） kl_div()（在 jax.scipy.special 模块中） kron()（在 jax.numpy 模块中）

L

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

logistic()（在 jax.lax 模块中） logistic()（在 jax.random 模块中） logit()（在 jax.scipy.special 模块中） lognormal()（在 jax.random 模块中） logpdf() （jax.scipy.stats.gaussian_kde 方法） logpdf()（在 jax.scipy.stats.beta 模块中） logpdf()（在 jax.scipy.stats.cauchy 模块中） logpdf()（在 jax.scipy.stats.chi2 模块中） logpdf()（在 jax.scipy.stats.dirichlet 模块中） logpdf()（在 jax.scipy.stats.expon 模块中） logpdf()（在 jax.scipy.stats.gamma 模块中） logpdf()（在 jax.scipy.stats.gennorm 模块中） logpdf()（在 jax.scipy.stats.laplace 模块中） logpdf()（在 jax.scipy.stats.logistic 模块中） logpdf()（在 jax.scipy.stats.multivariate_normal 模块中） logpdf()（在 jax.scipy.stats.norm 模块中） logpdf()（在 jax.scipy.stats.pareto 模块中） logpdf()（在 jax.scipy.stats.t 模块中） logpdf()（在 jax.scipy.stats.truncnorm 模块中） logpdf()（在 jax.scipy.stats.uniform 模块中） logpdf()（在 jax.scipy.stats.vonmises 模块中） logpdf()（在 jax.scipy.stats.wrapcauchy 模块中） logpmf()（在 jax.scipy.stats.bernoulli 模块中） logpmf()（在 jax.scipy.stats.betabinom 模块中） logpmf()（在 jax.scipy.stats.binom 模块中） logpmf()（在 jax.scipy.stats.geom 模块中） logpmf()（在 jax.scipy.stats.multinomial 模块中） logpmf()（在 jax.scipy.stats.nbinom 模块中） logpmf()（在 jax.scipy.stats.poisson 模块中） logsf()（在 jax.scipy.stats.beta 模块中） logsf()（在 jax.scipy.stats.cauchy 模块中） logsf()（在 jax.scipy.stats.chi2 模块中） logsf()（在 jax.scipy.stats.gamma 模块中） logsf()（在 jax.scipy.stats.norm 模块中） logsf()（在 jax.scipy.stats.truncnorm 模块中） logspace()（在 jax.numpy 模块中） logsumexp()（在 jax.nn 模块中） logsumexp()（在 jax.scipy.special 模块中） lower() （jax.stages.Wrapped 方法） Lowered（jax.stages 中的类） lpmn()（在 jax.scipy.special 模块中） lpmn_values()（在 jax.scipy.special 模块中） lstsq()（在 jax.numpy.linalg 模块中） lt()（在 jax.lax 模块中） lu()（在 jax.lax.linalg 模块中） lu()（在 jax.scipy.linalg 模块中） lu_factor()（在 jax.scipy.linalg 模块中） lu_solve()（在 jax.scipy.linalg 模块中）

M

machep（jax.numpy.finfo 属性） make_array_from_callback()（在 jax 模块中） make_array_from_process_local_data()（在 jax 模块中） make_array_from_single_device_arrays()（在 jax 模块中） make_jaxpr()（在 jax 模块中） make_schedule()（在 jax.example_libraries.optimizers 模块中） map()（在 jax.lax 模块中） map()（在 jax.tree 模块中） map_coordinates()（在 jax.scipy.ndimage 模块中） mask_indices()（在 jax.numpy 模块中） matmul()（在 jax.numpy 模块中） matmul()（在 jax.numpy.linalg 模块中） matrix_norm()（在 jax.numpy.linalg 模块中） matrix_power()（在 jax.numpy.linalg 模块中） matrix_rank()（在 jax.numpy.linalg 模块中） matrix_transpose()（在 jax.numpy 模块中） matrix_transpose()（在 jax.numpy.linalg 模块中） max() （jax.Array 方法） max()（在 jax.lax 模块中） max()（在 jax.numpy 模块中） maxexp（jax.numpy.finfo 属性） maximum()（在 jax.numpy 模块中） maximum_supported_calling_convention_version()（在 jax.export 模块中） MaxPool()（在 jax.example_libraries.stax 模块中） maxwell()（在 jax.random 模块中） max（jax.numpy.finfo 属性） mean() （jax.Array 方法） mean()（在 jax.numpy 模块中） median()（在 jax.numpy 模块中） memory_analysis() （jax.stages.Compiled 方法） memory_kind（jax.sharding.GSPMDSharding 属性） memory_kind（jax.sharding.NamedSharding 属性） memory_kind（jax.sharding.PmapSharding 属性） memory_kind（jax.sharding.PositionalSharding 属性） memory_kind（jax.sharding.Sharding 属性） memory_kind（jax.sharding.SingleDeviceSharding 属性） merge_linear_aux()（在 jax.extend.linear_util 模块中） meshgrid()（在 jax.numpy 模块中） Mesh（jax.sharding 中的类） mesh（jax.sharding.NamedSharding 属性） mgrid()（在 jax.numpy 模块中） min() （jax.Array 方法） min()（在 jax.lax 模块中） min()（在 jax.numpy 模块中） minexp（jax.numpy.finfo 属性） minimize()（在 jax.scipy.optimize 模块中） minimum()（在 jax.numpy 模块中） minimum_supported_calling_convention_version()（在 jax.export 模块中） min（jax.numpy.finfo 属性） mish()（在 jax.nn 模块中） mlir_module_serialized（jax.export.Exported 属性） mod()（在 jax.numpy 模块中） mode()（在 jax.scipy.stats 模块中） modf()（在 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.checkify jax.experimental.compilation_cache.compilation_cache jax.experimental.custom_partitioning jax.experimental.host_callback jax.experimental.jet jax.experimental.key_reuse jax.experimental.mesh_utils jax.experimental.multihost_utils jax.experimental.pallas jax.experimental.pjit jax.experimental.serialize_executable jax.experimental.shard_map jax.experimental.sparse jax.experimental.sparse.linalg jax.export jax.extend 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.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.tree jax.tree_util jax.typing

module_kept_var_idx（jax.export.Exported 属性） momentum()（在 jax.example_libraries.optimizers 模块中） moveaxis()（在 jax.numpy 模块中） mT（jax.Array 属性） mul()（在 jax.lax 模块中） multi_dot()（在 jax.numpy.linalg 模块中） multigammaln()（在 jax.scipy.special 模块中） multiply()（在 jax.numpy 模块中） multivariate_normal()（在 jax.random 模块中）

N

named_call()（在 jax 模块中） named_scope()（在 jax 模块中） NamedSharding（jax.sharding 中的类） nan_checks()（在 jax.experimental.checkify 模块中） nan_to_num()（在 jax.numpy 模块中） nanargmax()（在 jax.numpy 模块中） nanargmin()（在 jax.numpy 模块中） nancumprod()（在 jax.numpy 模块中） nancumsum()（在 jax.numpy 模块中） nanmax()（在 jax.numpy 模块中） nanmean()（在 jax.numpy 模块中） nanmedian()（在 jax.numpy 模块中） nanmin()（在 jax.numpy 模块中） nanpercentile()（在 jax.numpy 模块中） nanprod()（在 jax.numpy 模块中） nanquantile()（在 jax.numpy 模块中） nanstd()（在 jax.numpy 模块中） nansum()（在 jax.numpy 模块中） nanvar()（在 jax.numpy 模块中） nbytes（jax.Array 属性） ndarray()（在 jax.numpy 模块中） ndim()（在 jax.numpy 模块中） ndim（jax.Array 属性） ndtr()（在 jax.scipy.special 模块中） ndtri()（在 jax.scipy.special 模块中）

ne()（在 jax.lax 模块中） neg()（在 jax.lax 模块中） negative()（在 jax.numpy 模块中） negep（jax.numpy.finfo 属性） nesterov()（在 jax.example_libraries.optimizers 模块中） nexp（jax.numpy.finfo 属性） nextafter()（在 jax.lax 模块中） nextafter()（在 jax.numpy 模块中） nmant（jax.numpy.finfo 属性） NonConcreteBooleanIndexError（jax.errors 中的类） nonzero() （jax.Array 方法） nonzero()（在 jax.numpy 模块中） norm()（在 jax.numpy.linalg 模块中） normal()（在 jax.nn.initializers 模块中） normal()（在 jax.random 模块中） not_equal()（在 jax.numpy 模块中） nr_devices（jax.export.Exported 属性） num_devices（jax.sharding.GSPMDSharding 属性） num_devices（jax.sharding.NamedSharding 属性） num_devices（jax.sharding.PmapSharding 属性） num_devices（jax.sharding.PositionalSharding 属性） num_devices（jax.sharding.Sharding 属性） num_devices（jax.sharding.SingleDeviceSharding 属性） num_programs()（在 jax.experimental.pallas 模块中） number（jax.numpy 中的类） numpy_rank_promotion()（在 jax 模块中）

O

object_（jax.numpy 中的类） ogrid()（在 jax.numpy 模块中） one_hot()（在 jax.nn 模块中） ones()（在 jax.nn.initializers 模块中） ones()（在 jax.numpy 模块中） ones_like()（在 jax.numpy 模块中） optimization_barrier()（在 jax.lax 模块中） optimizer()（在 jax.example_libraries.optimizers 模块中） OptimizeResults（jax.scipy.optimize 中的类）

OptimizerState（jax.example_libraries.optimizers 中的类） Optimizer（jax.example_libraries.optimizers 中的类） ordered_effects（jax.export.Exported 属性） orthogonal()（在 jax.nn.initializers 模块中） orthogonal()（在 jax.random 模块中） out_avals（jax.export.Exported 属性） out_shardings_hlo（jax.export.Exported 属性） out_tree（jax.export.Exported 属性） outer()（在 jax.numpy 模块中） outer()（在 jax.numpy.linalg 模块中）

P

pack_optimizer_state()（在 jax.example_libraries.optimizers 模块中） packbits()（在 jax.numpy 模块中） packed_state（jax.example_libraries.optimizers.OptimizerState 属性） pad()（在 jax.lax 模块中） pad()（在 jax.numpy 模块中） pallas_call()（在 jax.experimental.pallas 模块中） parallel()（在 jax.example_libraries.stax 模块中） params_fn（jax.example_libraries.optimizers.Optimizer 属性） pareto()（在 jax.random 模块中） Partial（jax.tree_util 中的类） partition()（在 jax.numpy 模块中） PartitionSpec（jax.sharding 中的类） pdf() （jax.scipy.stats.gaussian_kde 方法） pdf()（在 jax.scipy.stats.beta 模块中） pdf()（在 jax.scipy.stats.cauchy 模块中） pdf()（在 jax.scipy.stats.chi2 模块中） pdf()（在 jax.scipy.stats.dirichlet 模块中） pdf()（在 jax.scipy.stats.expon 模块中） pdf()（在 jax.scipy.stats.gamma 模块中） pdf()（在 jax.scipy.stats.gennorm 模块中） pdf()（在 jax.scipy.stats.laplace 模块中） pdf()（在 jax.scipy.stats.logistic 模块中） pdf()（在 jax.scipy.stats.multivariate_normal 模块中） pdf()（在 jax.scipy.stats.norm 模块中） pdf()（在 jax.scipy.stats.pareto 模块中） pdf()（在 jax.scipy.stats.t 模块中） pdf()（在 jax.scipy.stats.truncnorm 模块中） pdf()（在 jax.scipy.stats.uniform 模块中） pdf()（在 jax.scipy.stats.vonmises 模块中） pdf()（在 jax.scipy.stats.wrapcauchy 模块中） percentile()（在 jax.numpy 模块中） permutation()（在 jax.random 模块中） permute_dims()（在 jax.numpy 模块中） piecewise()（在 jax.numpy 模块中） piecewise_constant()（在 jax.example_libraries.optimizers 模块中） pinv()（在 jax.numpy.linalg 模块中） pjit()（在 jax.experimental.pjit 模块中） place()（在 jax.numpy 模块中） platform_dependent()（在 jax.lax 模块中） platforms（jax.export.Exported 属性） pmap()（在 jax 模块中） PmapSharding（jax.sharding 中的类） pmax()（在 jax.lax 模块中） pmean()（在 jax.lax 模块中） pmf()（在 jax.scipy.stats.bernoulli 模块中） pmf()（在 jax.scipy.stats.betabinom 模块中） pmf()（在 jax.scipy.stats.binom 模块中） pmf()（在 jax.scipy.stats.geom 模块中） pmf()（在 jax.scipy.stats.multinomial 模块中） pmf()（在 jax.scipy.stats.nbinom 模块中） pmf()（在 jax.scipy.stats.poisson 模块中） pmin()（在 jax.lax 模块中）

poch()（在 jax.scipy.special 模块中） poisson()（在 jax.random 模块中） polar()（在 jax.scipy.linalg 模块中） poly()（在 jax.numpy 模块中） polyadd()（在 jax.numpy 模块中） polyder()（在 jax.numpy 模块中） polydiv()（在 jax.numpy 模块中） polyfit()（在 jax.numpy 模块中） polygamma()（在 jax.lax 模块中） polygamma()（在 jax.scipy.special 模块中） polyint()（在 jax.numpy 模块中） polymul()（在 jax.numpy 模块中） polynomial_decay()（在 jax.example_libraries.optimizers 模块中） polysub()（在 jax.numpy 模块中） polyval()（在 jax.numpy 模块中） population_count()（在 jax.lax 模块中） PositionalSharding（jax.sharding 中的类） positive()（在 jax.numpy 模块中） pow()（在 jax.lax 模块中） pow()（在 jax.numpy 模块中） power()（在 jax.numpy 模块中） ppermute()（在 jax.lax 模块中） ppf()（在 jax.scipy.stats.bernoulli 模块中） ppf()（在 jax.scipy.stats.cauchy 模块中） ppf()（在 jax.scipy.stats.logistic 模块中） ppf()（在 jax.scipy.stats.norm 模块中） ppf()（在 jax.scipy.stats.uniform 模块中） PrecisionLike()（在 jax.lax 模块中） Precision（jax.lax 中的类） precision（jax.numpy.finfo 属性） print()（在 jax.debug 模块中） print_environment_info()（在 jax 模块中） printoptions()（在 jax.numpy 模块中） prng_key（jax.dtypes 中的类） PRNGKey()（在 jax.random 模块中） process_allgather()（在 jax.experimental.multihost_utils 模块中） process_count()（在 jax 模块中） process_index()（在 jax 模块中） process_indices()（在 jax 模块中） prod() （jax.Array 方法） prod()（在 jax.numpy 模块中） program_id()（在 jax.experimental.pallas 模块中） promote_types()（在 jax.numpy 模块中） pshuffle()（在 jax.lax 模块中） psum()（在 jax.lax 模块中） psum_scatter()（在 jax.lax 模块中） pswapaxes()（在 jax.lax 模块中） ptp() （jax.Array 方法） ptp()（在 jax.numpy 模块中） pure_callback()（在 jax 模块中） put()（在 jax.numpy 模块中） pycapsule()（在 jax.extend.ffi 模块中） pytree

Q

qdwh()（在 jax.lax.linalg 模块中） qr()（在 jax.lax.linalg 模块中）

qr()（在 jax.numpy.linalg 模块中） qr()（在 jax.scipy.linalg 模块中） quantile()（在 jax.numpy 模块中）

R

r_()（在 jax.numpy 模块中） rad2deg()（在 jax.numpy 模块中） rademacher()（在 jax.random 模块中） radians()（在 jax.numpy 模块中） randint()（在 jax.random 模块中） random_bcoo()（在 jax.experimental.sparse 模块中） random_gamma_grad()（在 jax.lax 模块中） rankdata()（在 jax.scipy.stats 模块中） ravel() （jax.Array 方法） ravel()（在 jax.numpy 模块中） ravel_multi_index()（在 jax.numpy 模块中） ravel_pytree()（在 jax.flatten_util 模块中） rayleigh()（在 jax.random 模块中） rbg_prng_impl()（在 jax.extend.random 模块中） real()（在 jax.lax 模块中） real()（在 jax.numpy 模块中） real（jax.Array 属性） reciprocal()（在 jax.lax 模块中） reciprocal()（在 jax.numpy 模块中） reduce()（在 jax.lax 模块中） reduce()（在 jax.tree 模块中） reduce_precision()（在 jax.lax 模块中） reduce_window()（在 jax.lax 模块中） register_custom_call_target()（在 jaxlib.xla_client 模块中） register_dataclass()（在 jax.tree_util 模块中） register_ffi_target()（在 jax.extend.ffi 模块中） register_pytree_node()（在 jax.tree_util 模块中） register_pytree_node_class()（在 jax.tree_util 模块中） register_pytree_with_keys()（在 jax.tree_util 模块中） register_pytree_with_keys_class()（在 jax.tree_util 模块中） register_static()（在 jax.tree_util 模块中） RegularGridInterpolator（jax.scipy.interpolate 中的类） rel_entr()（在 jax.scipy.special 模块中） relu()（在 jax.nn 模块中） relu6()（在 jax.nn 模块中） rem()（在 jax.lax 模块中） remainder()（在 jax.numpy 模块中）

repeat() （jax.Array 方法） repeat()（在 jax.numpy 模块中） resample() （jax.scipy.stats.gaussian_kde 方法） reset_cache()（在 jax.experimental.compilation_cache.compilation_cache 模块中） reshape() （jax.Array 方法） reshape()（在 jax.lax 模块中） reshape()（在 jax.numpy 模块中） resize()（在 jax.image 模块中） resize()（在 jax.numpy 模块中） ResizeMethod（jax.image 中的类） resolution（jax.numpy.finfo 属性） result_type()（在 jax.dtypes 模块中） result_type()（在 jax.numpy 模块中） rev()（在 jax.lax 模块中） rfft()（在 jax.numpy.fft 模块中） rfft2()（在 jax.numpy.fft 模块中） rfftfreq()（在 jax.numpy.fft 模块中） rfftn()（在 jax.numpy.fft 模块中） right_shift()（在 jax.numpy 模块中） rint()（在 jax.numpy 模块中） rmsprop()（在 jax.example_libraries.optimizers 模块中） rmsprop_momentum()（在 jax.example_libraries.optimizers 模块中） rng_bit_generator()（在 jax.lax 模块中） rng_uniform()（在 jax.lax 模块中） roll()（在 jax.numpy 模块中） rollaxis()（在 jax.numpy 模块中） roots()（在 jax.numpy 模块中） rot90()（在 jax.numpy 模块中） Rotation（jax.scipy.spatial.transform 中的类） round() （jax.Array 方法） round()（在 jax.lax 模块中） round()（在 jax.numpy 模块中） round_()（在 jax.numpy 模块中） RoundingMethod（jax.lax 中的类） rsf2csf()（在 jax.scipy.linalg 模块中） rsqrt()（在 jax.lax 模块中） runtime_executable() （jax.stages.Compiled 方法）

S

s_()（在 jax.numpy 模块中） save()（在 jax.numpy 模块中） save_device_memory_profile()（在 jax.profiler 模块中） savez()（在 jax.numpy 模块中） scalar_type_of()（在 jax.dtypes 模块中） scale_and_translate()（在 jax.image 模块中） scan()（在 jax.lax 模块中） scatter()（在 jax.lax 模块中） scatter_add()（在 jax.lax 模块中） scatter_apply()（在 jax.lax 模块中） scatter_max()（在 jax.lax 模块中） scatter_min()（在 jax.lax 模块中） scatter_mul()（在 jax.lax 模块中） ScatterDimensionNumbers（jax.lax 中的类） schur()（在 jax.lax.linalg 模块中） schur()（在 jax.scipy.linalg 模块中） searchsorted() （jax.Array 方法） searchsorted()（在 jax.numpy 模块中） seed_with_impl()（在 jax.extend.random 模块中） segment_max()（在 jax.ops 模块中） segment_min()（在 jax.ops 模块中） segment_prod()（在 jax.ops 模块中） segment_sum()（在 jax.ops 模块中） select()（在 jax.lax 模块中） select()（在 jax.numpy 模块中） select_n()（在 jax.lax 模块中） selu()（在 jax.nn 模块中） sem()（在 jax.scipy.stats 模块中） serial()（在 jax.example_libraries.stax 模块中） serialize()（在 jax.experimental.serialize_executable 模块中） set_cache_dir()（在 jax.experimental.compilation_cache.compilation_cache 模块中） set_printoptions()（在 jax.numpy 模块中） setdiff1d()（在 jax.numpy 模块中） setxor1d()（在 jax.numpy 模块中） sf()（在 jax.scipy.stats.beta 模块中） sf()（在 jax.scipy.stats.cauchy 模块中） sf()（在 jax.scipy.stats.chi2 模块中） sf()（在 jax.scipy.stats.gamma 模块中） sf()（在 jax.scipy.stats.logistic 模块中） sf()（在 jax.scipy.stats.norm 模块中） sf()（在 jax.scipy.stats.truncnorm 模块中） sgd()（在 jax.example_libraries.optimizers 模块中） shape()（在 jax.numpy 模块中） shape_dependent()（在 jax.example_libraries.stax 模块中） ShapeDtypeStruct（jax 中的类） shape（jax.Array 属性） shard_map()（在 jax.experimental.shard_map 模块中） shard_shape() （jax.sharding.PmapSharding 方法） shard_shape() （jax.sharding.Sharding 方法） sharding_spec（jax.sharding.PmapSharding 属性） sharding（jax.Array 属性） Sharding（jax.sharding 中的类） shift_left()（在 jax.lax 模块中） shift_right_arithmetic()（在 jax.lax 模块中） shift_right_logical()（在 jax.lax 模块中） shutdown()（在 jax.distributed 模块中） sigmoid()（在 jax.nn 模块中） sign()（在 jax.lax 模块中） sign()（在 jax.numpy 模块中） signbit()（在 jax.numpy 模块中） signedinteger（jax.numpy 中的类） silu()（在 jax.nn 模块中） sin()（在 jax.lax 模块中） sin()（在 jax.numpy 模块中） sinc()（在 jax.numpy 模块中） single()（在 jax.numpy 模块中） SingleDeviceSharding（jax.sharding 中的类） sinh()（在 jax.lax 模块中） sinh()（在 jax.numpy 模块中） size()（在 jax.numpy 模块中） size（jax.Array 属性）

Slerp（jax.scipy.spatial.transform 中的类） slice()（在 jax.lax 模块中） slice_in_dim()（在 jax.lax 模块中） Slice（jax.experimental.pallas 中的类） slogdet()（在 jax.numpy.linalg 模块中） sm3()（在 jax.example_libraries.optimizers 模块中） smallest_normal（jax.numpy.finfo 属性） smallest_subnormal（jax.numpy.finfo 属性） soft_sign()（在 jax.nn 模块中） softmax()（在 jax.nn 模块中） softmax()（在 jax.scipy.special 模块中） softplus()（在 jax.nn 模块中） solve()（在 jax.numpy.linalg 模块中） solve()（在 jax.scipy.linalg 模块中） solve_triangular()（在 jax.scipy.linalg 模块中） sort() （jax.Array 方法） sort()（在 jax.lax 模块中） sort()（在 jax.numpy 模块中） sort_complex()（在 jax.numpy 模块中） sort_key_val()（在 jax.lax 模块中） sparse_plus()（在 jax.nn 模块中） sparse_sigmoid()（在 jax.nn 模块中） sparsify()（在 jax.experimental.sparse 模块中） spec（jax.sharding.NamedSharding 属性） spence()（在 jax.scipy.special 模块中） sph_harm()（在 jax.scipy.special 模块中） split()（在 jax.numpy 模块中） split()（在 jax.random 模块中） SPMD spsolve()（在 jax.experimental.sparse.linalg 模块中） sqrt()（在 jax.lax 模块中） sqrt()（在 jax.numpy 模块中） sqrtm()（在 jax.scipy.linalg 模块中） square()（在 jax.lax 模块中） square()（在 jax.numpy 模块中） squareplus()（在 jax.nn 模块中） squeeze() （jax.Array 方法） squeeze()（在 jax.lax 模块中） squeeze()（在 jax.numpy 模块中） stack()（在 jax.numpy 模块中） standardize()（在 jax.nn 模块中） start_server()（在 jax.profiler 模块中） start_trace()（在 jax.profiler 模块中） std() （jax.Array 方法） std()（在 jax.numpy 模块中） StepTraceAnnotation（jax.profiler 中的类） stft()（在 jax.scipy.signal 模块中） stop_gradient()（在 jax.lax 模块中） stop_trace()（在 jax.profiler 模块中） store()（在 jax.experimental.pallas 模块中） StoreException structure()（在 jax.tree 模块中） sub()（在 jax.lax 模块中） subtract()（在 jax.numpy 模块中） subtree_defs（jax.example_libraries.optimizers.OptimizerState 属性） sum() （jax.Array 方法） sum()（在 jax.numpy 模块中） SumPool()（在 jax.example_libraries.stax 模块中） svd()（在 jax.lax.linalg 模块中） svd()（在 jax.numpy.linalg 模块中） svd()（在 jax.scipy.linalg 模块中） svdvals()（在 jax.numpy.linalg 模块中） swap()（在 jax.experimental.pallas 模块中） swapaxes() （jax.Array 方法） swapaxes()（在 jax.numpy 模块中） swish()（在 jax.nn 模块中） switch()（在 jax.lax 模块中） symbolic_args_specs()（在 jax.export 模块中） symbolic_shape()（在 jax.export 模块中） SymbolicScope（jax.export 中的类） sync_global_devices()（在 jax.experimental.multihost_utils 模块中）

T

t()（在 jax.random 模块中） take() （jax.Array 方法） take()（在 jax.numpy 模块中） take_along_axis()（在 jax.numpy 模块中） tan()（在 jax.lax 模块中） tan()（在 jax.numpy 模块中） tanh()（在 jax.lax 模块中） tanh()（在 jax.numpy 模块中） tensordot()（在 jax.numpy 模块中） tensordot()（在 jax.numpy.linalg 模块中） tensorinv()（在 jax.numpy.linalg 模块中） tensorsolve()（在 jax.numpy.linalg 模块中） threefry2x32_p()（在 jax.extend.random 模块中） threefry_2x32()（在 jax.extend.random 模块中） threefry_prng_impl()（在 jax.extend.random 模块中） tile()（在 jax.numpy 模块中） tiny（jax.numpy.finfo 属性） to_device() （jax.Array 方法） to_dlpack()（在 jax.dlpack 模块中） todense()（在 jax.experimental.sparse 模块中） toeplitz()（在 jax.scipy.linalg 模块中） top_k()（在 jax.lax 模块中） TPU trace() （jax.Array 方法） trace() （jax.stages.Wrapped 方法） trace()（在 jax.numpy 模块中） trace()（在 jax.numpy.linalg 模块中） trace()（在 jax.profiler 模块中） TraceAnnotation（jax.profiler 中的类） TracerArrayConversionError（jax.errors 中的类） TracerBoolConversionError（jax.errors 中的类） TracerIntegerConversionError（jax.errors 中的类） transfer_guard()（在 jax 模块中） transformation()（在 jax.extend.linear_util 模块中） transformation_with_aux()（在 jax.extend.linear_util 模块中） transpose() （jax.Array 方法）

transpose()（在 jax.lax 模块中） transpose()（在 jax.numpy 模块中） transpose()（在 jax.tree 模块中） trapezoid()（在 jax.numpy 模块中） trapezoid()（在 jax.scipy.integrate 模块中） tree_all()（在 jax.tree_util 模块中） tree_def（jax.example_libraries.optimizers.OptimizerState 属性） tree_flatten()（在 jax.tree_util 模块中） tree_flatten_with_path()（在 jax.tree_util 模块中） tree_leaves()（在 jax.tree_util 模块中） tree_leaves_with_path()（在 jax.tree_util 模块中） tree_map()（在 jax.tree_util 模块中） tree_map_with_path()（在 jax.tree_util 模块中） tree_reduce()（在 jax.tree_util 模块中） tree_structure()（在 jax.tree_util 模块中） tree_transpose()（在 jax.tree_util 模块中） tree_unflatten()（在 jax.tree_util 模块中） treedef_children()（在 jax.tree_util 模块中） treedef_is_leaf()（在 jax.tree_util 模块中） treedef_tuple()（在 jax.tree_util 模块中） tri()（在 jax.numpy 模块中） triangular()（在 jax.random 模块中） triangular_solve()（在 jax.lax.linalg 模块中） tridiagonal()（在 jax.lax.linalg 模块中） tridiagonal_solve()（在 jax.lax.linalg 模块中） tril()（在 jax.numpy 模块中） tril_indices()（在 jax.numpy 模块中） tril_indices_from()（在 jax.numpy 模块中） trim_zeros()（在 jax.numpy 模块中） triu()（在 jax.numpy 模块中） triu_indices()（在 jax.numpy 模块中） triu_indices_from()（在 jax.numpy 模块中） true_divide()（在 jax.numpy 模块中） trunc()（在 jax.numpy 模块中） truncated_normal()（在 jax.nn.initializers 模块中） truncated_normal()（在 jax.random 模块中） T（jax.Array 属性）

U

ufunc（jax.numpy 中的类） uint()（在 jax.numpy 模块中） uint16（jax.numpy 中的类） uint32（jax.numpy 中的类） uint64（jax.numpy 中的类） uint8（jax.numpy 中的类） UnexpectedTracerError（jax.errors 中的类） unflatten()（在 jax.tree 模块中） uniform()（在 jax.nn.initializers 模块中） uniform()（在 jax.random 模块中） union1d()（在 jax.numpy 模块中） unique()（在 jax.numpy 模块中） unique_all()（在 jax.numpy 模块中）

unique_counts()（在 jax.numpy 模块中） unique_inverse()（在 jax.numpy 模块中） unique_values()（在 jax.numpy 模块中） unordered_effects（jax.export.Exported 属性） unpack_optimizer_state()（在 jax.example_libraries.optimizers 模块中） unpackbits()（在 jax.numpy 模块中） unravel_index()（在 jax.numpy 模块中） unsafe_rbg_prng_impl()（在 jax.extend.random 模块中） unsignedinteger（jax.numpy 中的类） unstack()（在 jax.numpy 模块中） unwrap()（在 jax.numpy 模块中） update_fn（jax.example_libraries.optimizers.Optimizer 属性） user_checks()（在 jax.experimental.checkify 模块中） uses_global_constants（jax.export.Exported 属性）

V

value_and_grad()（在 jax 模块中） value_and_grad()（在 jax.experimental.sparse 模块中） vander()（在 jax.numpy 模块中） var() （jax.Array 方法） var()（在 jax.numpy 模块中） variance_scaling()（在 jax.nn.initializers 模块中） vdot()（在 jax.numpy 模块中） vecdot()（在 jax.numpy 模块中） vecdot()（在 jax.numpy.linalg 模块中） vector_norm()（在 jax.numpy.linalg 模块中）

vectorize()（在 jax.numpy 模块中） view() （jax.Array 方法） visualize_array_sharding()（在 jax.debug 模块中） visualize_sharding()（在 jax.debug 模块中） VJP vjp()（在 jax 模块中） vmap()（在 jax 模块中） vq()（在 jax.scipy.cluster.vq 模块中） vsplit()（在 jax.numpy 模块中） vstack()（在 jax.numpy 模块中）

W

wald()（在 jax.random 模块中） weibull_min()（在 jax.random 模块中） welch()（在 jax.scipy.signal 模块中） where()（在 jax.numpy 模块中） while_loop()（在 jax.lax 模块中） with_memory_kind() （jax.sharding.GSPMDSharding 方法） with_memory_kind() （jax.sharding.NamedSharding 方法） with_memory_kind() （jax.sharding.PmapSharding 方法）

with_memory_kind() （jax.sharding.PositionalSharding 方法） with_memory_kind() （jax.sharding.Sharding 方法） with_memory_kind() （jax.sharding.SingleDeviceSharding 方法） with_sharding_constraint()（在 jax.lax 模块中） wrap_init()（在 jax.extend.linear_util 模块中） wrap_key_data()（在 jax.random 模块中） WrappedFun（jax.extend.linear_util 中的类） Wrapped（jax.stages 中的类）

X

XLA xla_computation()（在 jax 模块中）

xlog1py()（在 jax.scipy.special 模块中） xlogy()（在 jax.scipy.special 模块中）

Z

zeros()（在 jax.nn.initializers 模块中） zeros()（在 jax.numpy 模块中） zeros_like()（在 jax.numpy 模块中）

zeros_like_array()（在 jax.lax 模块中） zeta()（在 jax.lax 模块中） zeta()（在 jax.scipy.special 模块中）

函

函数式编程

前

前向模式自动微分

原

原始的

反

反向模式自动微分

变

变换

弱

弱类型

数

数组

纯

纯函数

设

设备

追

追踪器

静

静态