Toggle navigation

Index

_

__init__() (jax.core.ClosedJaxpr method) (jax.core.Jaxpr method) (jax.custom_jvp method) (jax.custom_vjp method) (jax.experimental.global_device_array.GlobalDeviceArray method) (jax.experimental.global_device_array.Shard method) (jax.experimental.sparse.BCOO method) (jax.numpy.bool_ method) (jax.numpy.character method) (jax.numpy.complex128 method) (jax.numpy.complex64 method) (jax.numpy.complexfloating method) (jax.numpy.dtype method) (jax.numpy.finfo method) (jax.numpy.flexible method) (jax.numpy.float16 method) (jax.numpy.float32 method) (jax.numpy.float64 method) (jax.numpy.floating method) (jax.numpy.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.ndarray method) (jax.numpy.number method) (jax.numpy.object_ method) (jax.numpy.signedinteger 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.StepTraceContext method) (jax.profiler.TraceAnnotation method) (jax.profiler.TraceContext method) (jax.scipy.optimize.OptimizeResults method) (jax.tree_util.Partial method) (jaxlib.xla_extension.CpuDevice method) (jaxlib.xla_extension.Device method) (jaxlib.xla_extension.GpuDevice method) (jaxlib.xla_extension.TpuDevice method)

A

abs() (in module jax.lax) (in module jax.numpy) absolute() (in module jax.numpy) acos() (in module jax.lax) adagrad() (in module jax.example_libraries.optimizers) adam() (in module jax.example_libraries.optimizers) adamax() (in module jax.example_libraries.optimizers) add() (in module jax.lax) (in module jax.numpy) all() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) 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) alltrue() (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) (jaxlib.xla_extension.DeviceArray method) append() (in module jax.numpy) apply_along_axis() (in module jax.numpy) apply_over_axes() (in module jax.numpy) approx_max_k() (in module jax.experimental.ann) approx_min_k() (in module jax.experimental.ann) 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) (jaxlib.xla_extension.DeviceArray method) argmin() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) argpartition() (jaxlib.xla_extension.DeviceArray method) argsort() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) argwhere() (in module jax.numpy) around() (in module jax.numpy) 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) asarray() (in module jax.numpy) asin() (in module jax.lax) associative_scan() (in module jax.lax) at (jax.numpy.ndarray property) (jaxlib.xla_extension.DeviceArray property) atan() (in module jax.lax) atan2() (in module jax.lax) atleast_1d() (in module jax.numpy) atleast_2d() (in module jax.numpy) atleast_3d() (in module jax.numpy) average() (in module jax.numpy) AvgPool() (in module jax.example_libraries.stax) axis_index() (in module jax.lax)

B

barrier_wait() (in module jax.experimental.host_callback) bartlett() (in module jax.numpy) batch_matmul() (in module jax.lax) BatchNorm() (in module jax.example_libraries.stax) BCOO (class in jax.experimental.sparse) bernoulli() (in module jax.random) bessel_i0e() (in module jax.lax) bessel_i1e() (in module jax.lax) beta() (in module jax.random) betainc() (in module jax.lax) (in module jax.scipy.special) bicgstab() (in module jax.scipy.sparse.linalg) bincount() (in module jax.numpy) bitcast_convert_type() (in module jax.lax) bits (jax.numpy.finfo attribute) (jax.numpy.iinfo attribute) bitwise_and() (in module jax.lax) (in module jax.numpy) bitwise_not() (in module jax.lax) (in module jax.numpy)

bitwise_or() (in module jax.lax) (in module jax.numpy) bitwise_xor() (in module jax.lax) (in module jax.numpy) blackman() (in module jax.numpy) block() (in module jax.numpy) block_diag() (in module jax.scipy.linalg) block_host_until_ready() (jaxlib.xla_extension.DeviceArray method) block_until_ready() (in module jax) (jaxlib.xla_extension.DeviceArray method) bool_ (class in jax.numpy) broadcast() (in module jax.lax) (jaxlib.xla_extension.DeviceArray method) broadcast_arrays() (in module jax.numpy) broadcast_in_dim() (in module jax.lax) (jaxlib.xla_extension.DeviceArray method) broadcast_shapes() (in module jax.numpy) broadcast_to() (in module jax.numpy) broadcasted_iota() (in module jax.lax) build_tree() (in module jax.tree_util)

C

c_ (in module jax.numpy) call() (in module jax.experimental.host_callback) CallbackException can_cast() (in module jax.numpy) 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.laplace) (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) cdouble (in module jax.numpy) ceil() (in module jax.lax) (in module jax.numpy) celu() (in module jax.nn) cg() (in module jax.scipy.sparse.linalg) character (class in jax.numpy) check() (in module jax.experimental.checkify) check_error() (in module jax.experimental.checkify) check_tracer_leaks (in module jax) checkify() (in module jax.experimental.checkify) checking_leaks (in module jax) checkpoint() (in module jax) cho_factor() (in module jax.scipy.linalg) cho_solve() (in module jax.scipy.linalg) choice() (in module jax.random) cholesky() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) choose() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) clamp() (in module jax.lax) clip() (in module jax.numpy) clip_grads() (in module jax.example_libraries.optimizers) clone() (jaxlib.xla_extension.DeviceArray method) ClosedJaxpr (class in jax.core) closure_convert() (in module jax) collapse() (in module jax.lax) column_stack() (in module jax.numpy) complex() (in module jax.lax) complex128 (class in jax.numpy) complex64 (class in jax.numpy) complex_ (in module jax.numpy) complexfloating (class in jax.numpy) ComplexWarning compress() (in module jax.numpy) concatenate() (in module jax.lax) (in module jax.numpy) ConcretizationTypeError (class in jax.errors) cond() (in module jax.lax) (in module jax.numpy.linalg) cond_range() (jax.experimental.loops.Scope method) config (in module jax)

conj() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) conjugate() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) constant() (in module jax.example_libraries.optimizers) 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) copy() (jaxlib.xla_extension.DeviceArray method) copy_to_device() (jaxlib.xla_extension.DeviceArray method) copy_to_host_async() (jaxlib.xla_extension.DeviceArray 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) count_nonzero() (in module jax.numpy) cov() (in module jax.numpy) CPU CpuDevice (class in jaxlib.xla_extension) cross() (in module jax.numpy) csd() (in module jax.scipy.signal) csingle (in module jax.numpy) cummax() (in module jax.lax) cummin() (in module jax.lax) cumprod() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) cumproduct() (in module jax.numpy) cumsum() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) custom_jvp (class in jax) custom_linear_solve() (in module jax.lax) custom_root() (in module jax.lax) custom_vjp (class in jax)

D

data (jax.experimental.sparse.BCOO attribute) dct() (in module jax.scipy.fft) dctn() (in module jax.scipy.fft) debug_infs (in module jax) debug_nans (in module jax) default_backend() (in module jax) (in module jax.lib.xla_bridge) default_matmul_precision (in module jax) default_prng_impl (in module jax) deg2rad() (in module jax.numpy) degrees() (in module jax.numpy) delete() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) Dense() (in module jax.example_libraries.stax) det (in module jax.numpy.linalg) det() (in module jax.scipy.linalg) Device (class in jaxlib.xla_extension) device() (jaxlib.xla_extension.DeviceArray method) device_count() (in module jax) (in module jax.lib.xla_bridge) device_get() (in module jax) device_memory_profile() (in module jax.profiler) device_put() (in module jax) device_put_replicated() (in module jax) device_put_sharded() (in module jax) DeviceArray (class in jaxlib.xla_extension) (in module jax.numpy) DeviceArrayBase (class in jaxlib.xla_extension) devices() (in module jax) 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) (jaxlib.xla_extension.DeviceArray method) diff() (in module jax.numpy) digamma() (in module jax.lax) (in module jax.scipy.special) digitize() (in module jax.numpy) dirichlet() (in module jax.random) disable_jit() (in module jax) disable_x64() (in module jax.experimental) div() (in module jax.lax) divide() (in module jax.numpy) divmod() (in module jax.numpy) dot() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) dot_general() (in module jax.lax) double (in module jax.numpy) double_sided_maxwell() (in module jax.random) Dropout() (in module jax.example_libraries.stax) dsplit() (in module jax.numpy) dstack() (in module jax.numpy) dtype (class in jax.numpy) (jax.experimental.global_device_array.GlobalDeviceArray attribute) 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) eigvals() (in module jax.numpy.linalg) eigvalsh() (in module jax.numpy.linalg) einsum() (in module jax.numpy) einsum_path() (in module jax.numpy) elementwise() (in module jax.example_libraries.stax) elu() (in module jax.nn) empty() (in module jax.numpy) empty_like() (in module jax.numpy) enable_checks (in module jax) enable_custom_prng (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) eval_shape() (in module jax) exp() (in module jax.lax) (in module jax.numpy) exp1() (in module jax.scipy.special) exp2() (in module jax.numpy) expand_dims() (in module jax.lax) (in module jax.numpy) expi (in module jax.scipy.special) expit (in module jax.scipy.special) expm() (in module jax.scipy.linalg) expm1() (in module jax.lax) (in module jax.numpy) expm_frechet() (in module jax.scipy.linalg) expn (in module jax.scipy.special) exponential() (in module jax.random) exponential_decay() (in module jax.example_libraries.optimizers) extract() (in module jax.numpy) eye() (in module jax.numpy)

F

fabs() (in module jax.numpy) FanInConcat() (in module jax.example_libraries.stax) FanOut() (in module jax.example_libraries.stax) fft() (in module jax.lax) (in module jax.numpy.fft) fft2() (in module jax.numpy.fft) fftfreq() (in module jax.numpy.fft) fftn() (in module jax.numpy.fft) fftshift() (in module jax.numpy.fft) finfo (class in jax.numpy) fix() (in module jax.numpy) flatnonzero() (in module jax.numpy) flatten() (jaxlib.xla_extension.DeviceArray method) flexible (class in jax.numpy) flip() (in module jax.numpy) fliplr() (in module jax.numpy) flipud() (in module jax.numpy) float16 (class in jax.numpy) float32 (class in jax.numpy)

float64 (class in jax.numpy) float_ (in module jax.numpy) float_power() (in module jax.numpy) floating (class in jax.numpy) floor() (in module jax.lax) (in module jax.numpy) floor_divide() (in module jax.numpy) fmax() (in module jax.numpy) fmin() (in module jax.numpy) fmod() (in module jax.numpy) fold_in() (in module jax.random) fori_loop() (in module jax.lax) forward-mode autodiff frexp() (in module jax.numpy) from_dlpack() (in module jax.dlpack) full() (in module jax.lax) (in module jax.numpy) full_like() (in module jax.lax) (in module jax.numpy) functional programming

G

gamma() (in module jax.random) gammainc() (in module jax.scipy.special) gammaincc() (in module jax.scipy.special) gammaln() (in module jax.scipy.special) gather() (in module jax.lax) GatherDimensionNumbers (class in jax.lax) GatherScatterMode (class in jax.lax) 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) geomspace() (in module jax.numpy) get_backend() (in module jax.lib.xla_bridge) get_compile_options() (in module jax.lib.xla_bridge)

get_printoptions() (in module jax.numpy) global_shards (jax.experimental.global_device_array.GlobalDeviceArray attribute) GlobalDeviceArray (class in jax.experimental.global_device_array) glorot_normal() (in module jax.nn.initializers) glorot_uniform() (in module jax.nn.initializers) glu() (in module jax.nn) gmres() (in module jax.scipy.sparse.linalg) GPU GpuDevice (class in jaxlib.xla_extension) grad() (in module jax) gradient() (in module jax.numpy) greater() (in module jax.numpy) greater_equal() (in module jax.numpy) gt() (in module jax.lax) gumbel() (in module jax.random)

H

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

hessian() (in module jax) hfft() (in module jax.numpy.fft) histogram() (in module jax.numpy) histogram2d() (in module jax.numpy) histogram_bin_edges() (in module jax.numpy) histogramdd() (in module jax.numpy) hsplit() (in module jax.numpy) hstack() (in module jax.numpy) 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) id_print() (in module jax.experimental.host_callback) id_tap() (in module jax.experimental.host_callback) identity() (in module jax.numpy) iexp (jax.numpy.finfo attribute) ifft() (in module jax.numpy.fft) ifft2() (in module jax.numpy.fft) ifftn() (in module jax.numpy.fft) ifftshift() (in module jax.numpy.fft) igamma() (in module jax.lax) igammac() (in module jax.lax) ihfft() (in module jax.numpy.fft) iinfo (class in jax.numpy) imag (jaxlib.xla_extension.DeviceArray property) imag() (in module jax.lax) (in module jax.numpy) in1d() (in module jax.numpy) index (in module jax.ops) index_add() (in module jax.ops) index_exp (in module jax.numpy) index_in_dim() (in module jax.lax) index_max() (in module jax.ops) index_min() (in module jax.ops) index_mul() (in module jax.ops) index_take() (in module jax.lax) index_update() (in module jax.ops) 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) inner() (in module jax.numpy)

insert() (in module jax.numpy) 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) 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) iota() (in module jax.lax) irfft() (in module jax.numpy.fft) irfft2() (in module jax.numpy.fft) irfftn() (in module jax.numpy.fft) is_deleted() (jaxlib.xla_extension.DeviceArray method) is_finite() (in module jax.lax) is_fully_replicated (jax.experimental.global_device_array.GlobalDeviceArray attribute) isclose() (in module jax.numpy) iscomplex() (in module jax.numpy) iscomplexobj() (in module jax.numpy) isf() (in module jax.scipy.stats.logistic) 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.numpy) issubsctype() (in module jax.numpy) iterable() (in module jax.numpy) ix_() (in module jax.numpy)

J

jacfwd() (in module jax) jacrev() (in module jax) jax.core module jax.distributed module jax.dlpack module jax.example_libraries module jax.example_libraries.optimizers module jax.example_libraries.stax module jax.experimental.ann module jax.experimental.global_device_array module jax.experimental.host_callback module jax.experimental.loops module jax.experimental.maps module jax.experimental.pjit module jax.experimental.sparse module jax.flatten_util module jax.image module jax.lax module jax.lax.linalg module jax.nn module jax.nn.initializers module jax.numpy module jax.numpy.fft module jax.numpy.linalg module jax.ops module jax.profiler module jax.random module jax.scipy.fft module

jax.scipy.linalg module jax.scipy.ndimage module jax.scipy.optimize module jax.scipy.signal module jax.scipy.sparse.linalg module jax.scipy.special module jax.scipy.stats.bernoulli module jax.scipy.stats.beta module jax.scipy.stats.betabinom 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.geom module jax.scipy.stats.laplace module jax.scipy.stats.logistic module jax.scipy.stats.multivariate_normal module jax.scipy.stats.norm module jax.scipy.stats.pareto module jax.scipy.stats.poisson module jax.scipy.stats.t module jax.scipy.stats.uniform module jax.tree_util module jaxpr Jaxpr (class in jax.core) JIT jit() (in module jax) JoinPoint (class in jax.example_libraries.optimizers) JVP jvp() (in module jax)

K

kaiser() (in module jax.numpy)

kron() (in module jax.numpy)

L

l2_norm() (in module jax.example_libraries.optimizers) laplace() (in module jax.random) lcm() (in module jax.numpy) ldexp() (in module jax.numpy) le() (in module jax.lax) leaky_relu() (in module jax.nn) lecun_normal() (in module jax.nn.initializers) lecun_uniform() (in module jax.nn.initializers) left_shift() (in module jax.numpy) less() (in module jax.numpy) less_equal() (in module jax.numpy) lexsort() (in module jax.numpy) lgamma() (in module jax.lax) linear_transpose() (in module jax) linearize() (in module jax) linspace() (in module jax.numpy) load() (in module jax.numpy) local_device_count() (in module jax) (in module jax.lib.xla_bridge) local_devices() (in module jax) local_shards (jax.experimental.global_device_array.GlobalDeviceArray attribute) log() (in module jax.lax) (in module jax.numpy) log10() (in module jax.numpy) log1p() (in module jax.lax) (in module jax.numpy) log2() (in module jax.numpy) log_compiles (in module jax) log_ndtr (in module jax.scipy.special) log_sigmoid() (in module jax.nn) log_softmax() (in module jax.nn) logaddexp (in module jax.numpy) logaddexp2 (in module jax.numpy) logcdf() (in module jax.scipy.stats.norm)

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.random) logit (in module jax.scipy.special) 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.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.uniform) logpmf() (in module jax.scipy.stats.bernoulli) (in module jax.scipy.stats.betabinom) (in module jax.scipy.stats.geom) (in module jax.scipy.stats.poisson) logspace() (in module jax.numpy) logsumexp() (in module jax.nn) (in module jax.scipy.special) lpmn() (in module jax.scipy.special) lpmn_values() (in module jax.scipy.special) lstsq() (in module jax.numpy.linalg) lt() (in module jax.lax) lu() (in module jax.lax.linalg) (in module jax.scipy.linalg) lu_factor() (in module jax.scipy.linalg) lu_solve() (in module jax.scipy.linalg)

M

machar (jax.numpy.finfo attribute) machep (jax.numpy.finfo attribute) make_jaxpr() (in module jax) make_schedule() (in module jax.example_libraries.optimizers) map() (in module jax.lax) map_coordinates() (in module jax.scipy.ndimage) mask_indices() (in module jax.numpy) matmul() (in module jax.numpy) matrix_power() (in module jax.numpy.linalg) matrix_rank() (in module jax.numpy.linalg) max (jax.numpy.finfo attribute) (jax.numpy.iinfo attribute) max() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) maxexp (jax.numpy.finfo attribute) maximum() (in module jax.numpy) MaxPool() (in module jax.example_libraries.stax) maxwell() (in module jax.random) mean() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) median() (in module jax.numpy) mesh() (in module jax.experimental.maps) meshgrid() (in module jax.numpy) mgrid (in module jax.numpy) min (jax.numpy.finfo attribute) (jax.numpy.iinfo attribute) min() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) minexp (jax.numpy.finfo attribute) minimize() (in module jax.scipy.optimize) minimum() (in module jax.numpy) mod() (in module jax.numpy) modf() (in module jax.numpy) module jax.core jax.distributed jax.dlpack jax.example_libraries jax.example_libraries.optimizers jax.example_libraries.stax jax.experimental.ann jax.experimental.global_device_array jax.experimental.host_callback jax.experimental.loops jax.experimental.maps jax.experimental.pjit jax.experimental.sparse jax.flatten_util jax.image jax.lax jax.lax.linalg jax.nn jax.nn.initializers jax.numpy jax.numpy.fft jax.numpy.linalg jax.ops jax.profiler jax.random jax.scipy.fft jax.scipy.linalg jax.scipy.ndimage jax.scipy.optimize jax.scipy.signal jax.scipy.sparse.linalg jax.scipy.special jax.scipy.stats.bernoulli jax.scipy.stats.beta jax.scipy.stats.betabinom jax.scipy.stats.cauchy jax.scipy.stats.chi2 jax.scipy.stats.dirichlet jax.scipy.stats.expon jax.scipy.stats.gamma jax.scipy.stats.geom jax.scipy.stats.laplace jax.scipy.stats.logistic jax.scipy.stats.multivariate_normal jax.scipy.stats.norm jax.scipy.stats.pareto jax.scipy.stats.poisson jax.scipy.stats.t jax.scipy.stats.uniform jax.tree_util

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

N

named_call() (in module jax) nan_to_num() (in module jax.numpy) nanargmax() (in module jax.numpy) nanargmin() (in module jax.numpy) nancumprod() (in module jax.numpy) nancumsum() (in module jax.numpy) nanmax() (in module jax.numpy) nanmean() (in module jax.numpy) nanmedian() (in module jax.numpy) nanmin() (in module jax.numpy) nanpercentile() (in module jax.numpy) nanprod() (in module jax.numpy) nanquantile() (in module jax.numpy) nanstd() (in module jax.numpy) nansum() (in module jax.numpy) nanvar() (in module jax.numpy) ndarray (class in jax.numpy) ndim (jax.experimental.global_device_array.GlobalDeviceArray attribute) ndim() (in module jax.numpy) ndtr() (in module jax.scipy.special)

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

O

object_ (class in jax.numpy) ogrid (in module jax.numpy) on_device_size_in_bytes() (jaxlib.xla_extension.DeviceArray method) one_hot() (in module jax.nn) ones() (in module jax.nn.initializers) (in module jax.numpy)

ones_like() (in module jax.numpy) Optimizer (class in jax.example_libraries.optimizers) optimizer() (in module jax.example_libraries.optimizers) OptimizeResults (class in jax.scipy.optimize) OptimizerState (class in jax.example_libraries.optimizers) outer() (in module jax.numpy)

P

pack_optimizer_state() (in module jax.example_libraries.optimizers) packbits() (in module jax.numpy) packed_state (jax.example_libraries.optimizers.OptimizerState attribute) pad() (in module jax.lax) (in module jax.numpy) parallel() (in module jax.example_libraries.stax) params_fn (jax.example_libraries.optimizers.Optimizer attribute) pareto() (in module jax.random) Partial (class in jax.tree_util) 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.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.uniform) percentile() (in module jax.numpy) permutation() (in module jax.random) piecewise() (in module jax.numpy) piecewise_constant() (in module jax.example_libraries.optimizers) pinv (in module jax.numpy.linalg) pjit() (in module jax.experimental.pjit) platform() (jaxlib.xla_extension.DeviceArray method) pmap() (in module jax) 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.geom) (in module jax.scipy.stats.poisson)

pmin() (in module jax.lax) poisson() (in module jax.random) poly() (in module jax.numpy) polyadd() (in module jax.numpy) polyder() (in module jax.numpy) polyfit() (in module jax.numpy) polygamma() (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) power() (in module jax.numpy) ppermute() (in module jax.lax) ppf() (in module jax.scipy.stats.logistic) (in module jax.scipy.stats.norm) Precision (class in jax.lax) precision (jax.numpy.finfo attribute) printoptions() (in module jax.numpy) PRNGKey() (in module jax.random) process_count() (in module jax) process_index() (in module jax) (in module jax.lib.xla_bridge) prod() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) product() (in module jax.numpy) promote_types() (in module jax.numpy) pshuffle() (in module jax.lax) psum() (in module jax.lax) pswapaxes() (in module jax.lax) ptp() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) pure function

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) range() (jax.experimental.loops.Scope method) ravel() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) ravel_multi_index() (in module jax.numpy) ravel_pytree() (in module jax.flatten_util) real (jaxlib.xla_extension.DeviceArray property) real() (in module jax.lax) (in module jax.numpy) reciprocal() (in module jax.lax) (in module jax.numpy) reduce() (in module jax.lax) reduce_precision() (in module jax.lax) reduce_window() (in module jax.lax) register_pytree_node() (in module jax.tree_util) register_pytree_node_class() (in module jax.tree_util) 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) (jaxlib.xla_extension.DeviceArray method)

reshape() (in module jax.lax) (in module jax.numpy) resize() (in module jax.image) (in module jax.numpy) resolution (jax.numpy.finfo attribute) result_type() (in module jax.numpy) rev() (in module jax.lax) reverse-mode autodiff rfft() (in module jax.numpy.fft) rfft2() (in module jax.numpy.fft) rfftfreq() (in module jax.numpy.fft) rfftn() (in module jax.numpy.fft) right_shift() (in module jax.numpy) rint() (in module jax.numpy) rmsprop() (in module jax.example_libraries.optimizers) rmsprop_momentum() (in module jax.example_libraries.optimizers) roll() (in module jax.numpy) rollaxis() (in module jax.numpy) roots() (in module jax.numpy) rot90() (in module jax.numpy) round() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) round_() (in module jax.numpy) RoundingMethod (class in jax.lax) row_stack() (in module jax.numpy) rsqrt() (in module jax.lax)

S

s_ (in module jax.numpy) save() (in module jax.numpy) save_device_memory_profile() (in module jax.profiler) savez() (in module jax.numpy) scale_and_translate() (in module jax.image) scan() (in module jax.lax) scatter() (in module jax.lax) scatter_add() (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.scipy.linalg) Scope (class in jax.experimental.loops) searchsorted() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) 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) selu() (in module jax.nn) serial() (in module jax.example_libraries.stax) set_printoptions() (in module jax.numpy) setdiff1d() (in module jax.numpy) setxor1d() (in module jax.numpy) sf() (in module jax.scipy.stats.logistic) sgd() (in module jax.example_libraries.optimizers) shape (jax.experimental.global_device_array.GlobalDeviceArray attribute) shape() (in module jax.numpy) shape_dependent() (in module jax.example_libraries.stax) Shard (class in jax.experimental.global_device_array) shift_left() (in module jax.lax) shift_right_arithmetic() (in module jax.lax) shift_right_logical() (in module jax.lax) shuffle() (in module jax.random) 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) sinh() (in module jax.lax) (in module jax.numpy) size (jax.experimental.global_device_array.GlobalDeviceArray attribute) size() (in module jax.numpy) slice() (in module jax.lax) slice_in_dim() (in module jax.lax) slogdet (in module jax.numpy.linalg)

sm3() (in module jax.example_libraries.optimizers) smallest_normal (jax.numpy.finfo attribute) smallest_subnormal (jax.numpy.finfo attribute) soft_sign() (in module jax.nn) softmax() (in module jax.nn) softplus() (in module jax.nn) solve() (in module jax.numpy.linalg) (in module jax.scipy.linalg) solve_triangular() (in module jax.scipy.linalg) sometrue() (in module jax.numpy) sort() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) sort_complex() (in module jax.numpy) sort_key_val() (in module jax.lax) sparsify() (in module jax.experimental.sparse) sph_harm() (in module jax.scipy.special) split() (in module jax.numpy) (in module jax.random) (jaxlib.xla_extension.DeviceArray method) SPMD sqrt() (in module jax.lax) (in module jax.numpy) sqrtm() (in module jax.scipy.linalg) square() (in module jax.lax) (in module jax.numpy) squeeze() (in module jax.lax) (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) stack() (in module jax.numpy) start_server() (in module jax.profiler) start_subtrace() (jax.experimental.loops.Scope method) start_trace() (in module jax.profiler) static std() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) StepTraceAnnotation (class in jax.profiler) StepTraceContext (class in jax.profiler) stft() (in module jax.scipy.signal) stop_gradient() (in module jax.lax) stop_trace() (in module jax.profiler) sub() (in module jax.lax) subtract() (in module jax.numpy) subtree_defs (jax.example_libraries.optimizers.OptimizerState attribute) sum() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) SumPool() (in module jax.example_libraries.stax) svd() (in module jax.lax.linalg) (in module jax.numpy.linalg) (in module jax.scipy.linalg) swapaxes() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) swish() (in module jax.nn) switch() (in module jax.lax)

T

T (jaxlib.xla_extension.DeviceArray property) t() (in module jax.random) take() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) take_along_axis() (in module jax.numpy) tan() (in module jax.lax) (in module jax.numpy) tanh() (in module jax.numpy) tensordot() (in module jax.numpy) tensorinv() (in module jax.numpy.linalg) tensorsolve() (in module jax.numpy.linalg) tie_in() (in module jax.lax) tile() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) tiny (jax.numpy.finfo attribute) to_dlpack() (in module jax.dlpack) to_py() (jaxlib.xla_extension.DeviceArray method) top_k() (in module jax.lax) TPU TpuDevice (class in jaxlib.xla_extension) trace() (in module jax.numpy) (in module jax.profiler) (jaxlib.xla_extension.DeviceArray method) trace_function() (in module jax.profiler) TraceAnnotation (class in jax.profiler) TraceContext (class in jax.profiler) Tracer TracerArrayConversionError (class in jax.errors) TracerIntegerConversionError (class in jax.errors)

transformation transpose() (in module jax.lax) (in module jax.numpy) trapz() (in module jax.numpy) tree_all() (in module jax.tree_util) tree_def (jax.example_libraries.optimizers.OptimizerState attribute) tree_flatten() (in module jax.tree_util) tree_leaves() (in module jax.tree_util) tree_map() (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_solve() (in module jax.lax.linalg) tril() (in module jax.numpy) (in module jax.scipy.linalg) tril_indices() (in module jax.numpy) tril_indices_from() (in module jax.numpy) trim_zeros() (in module jax.numpy) triu() (in module jax.numpy) (in module jax.scipy.linalg) triu_indices() (in module jax.numpy) triu_indices_from() (in module jax.numpy) true_divide() (in module jax.numpy) trunc() (in module jax.numpy) truncated_normal() (in module jax.random)

U

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

unique() (in module jax.numpy) unpack_optimizer_state() (in module jax.example_libraries.optimizers) unpackbits() (in module jax.numpy) unravel_index() (in module jax.numpy) unsafe_buffer_pointer() (jaxlib.xla_extension.DeviceArray method) unsignedinteger (class in jax.numpy) unwrap() (in module jax.numpy) update_fn (jax.example_libraries.optimizers.Optimizer attribute)

V

value_and_grad() (in module jax) vander() (in module jax.numpy) var() (in module jax.numpy) (jaxlib.xla_extension.DeviceArray method) variance_scaling() (in module jax.nn.initializers) vdot() (in module jax.numpy)

vectorize() (in module jax.numpy) VJP vjp() (in module jax) vmap() (in module jax) vsplit() (in module jax.numpy) vstack() (in module jax.numpy)

W

weak type weibull_min() (in module jax.random) welch() (in module jax.scipy.signal)

where() (in module jax.numpy) while_loop() (in module jax.lax) while_range() (jax.experimental.loops.Scope method)

X

XLA xla_computation() (in module jax) xla_dynamic_shape() (jaxlib.xla_extension.DeviceArray method)

xla_shape() (jaxlib.xla_extension.DeviceArray method) xlog1py() (in module jax.scipy.special) xlogy() (in module jax.scipy.special) xmap() (in module jax.experimental.maps)

Z

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

zeros_like() (in module jax.numpy) zeta() (in module jax.scipy.special)