René's URL Explorer Experiment


Title: JAX: High performance array computing — JAX documentation

direct link

Domain: jax.dev

docsearch:languageen
readthedocs-project-slugjax
readthedocs-version-sluglatest
readthedocs-resolver-filename/
readthedocs-http-status200

Links:

Skip to main contenthttps://jax.dev#main-content
https://jax.dev
Installationhttps://jax.dev/installation.html
Newly documentedhttps://jax.dev/whats-new.html
JAX 101: expressing computationshttps://jax.dev/101/index.html
Arrays and jax.numpyhttps://jax.dev/101/arrays.html
Transformations: grad and vmaphttps://jax.dev/101/transformations.html
Pytreeshttps://jax.dev/101/pytrees.html
Pseudorandom numbershttps://jax.dev/101/random.html
Stateful computationshttps://jax.dev/101/state.html
Errorshttps://jax.dev/101/errors.html
Convolutionshttps://jax.dev/101/convolutions.html
Default dtypes and the X64 flaghttps://jax.dev/101/default_dtypes.html
Type promotion semanticshttps://jax.dev/101/type_promotion.html
Rank promotion warninghttps://jax.dev/101/rank_promotion_warning.html
JAX 201: performance and scalinghttps://jax.dev/201/index.html
Just-in-time compilationhttps://jax.dev/201/jit.html
Ahead-of-time lowering and compilationhttps://jax.dev/201/aot.html
Control flow and logical operators with jithttps://jax.dev/201/control-flow.html
Data placementhttps://jax.dev/201/placement.html
Distributed arrays and automatic parallelizationhttps://jax.dev/201/sharding.html
Manual parallelism with shard_maphttps://jax.dev/201/shard-map.html
External callbackshttps://jax.dev/201/callbacks.html
Benchmarking and profilinghttps://jax.dev/201/profiling.html
Debugging runtime valueshttps://jax.dev/201/debugging.html
Debugging slow JAX tracing and XLA compilationhttps://jax.dev/201/slow-compilation.html
Matmul precisionhttps://jax.dev/201/precision.html
Controlling XLA from JAXhttps://jax.dev/201/controlling-xla.html
GPU memory allocationhttps://jax.dev/201/gpu-memory.html
Memory spaces and host offloadinghttps://jax.dev/201/memory-spaces.html
JAX 301: advanced autodiff and extending JAXhttps://jax.dev/301/index.html
The autodiff cookbook: JVPs and VJPshttps://jax.dev/301/cookbook.html
First-class VJPshttps://jax.dev/301/vjp-objects.html
Autodiff and shardinghttps://jax.dev/301/sharding-ad.html
Custom derivative rules with hijax primitiveshttps://jax.dev/301/custom-derivatives.html
Custom JVPs or VJPs with custom_jvp and custom_vjphttps://jax.dev/301/custom-jvp-vjp.html
Autodiff with refshttps://jax.dev/301/refs.html
Gradient checkpointing with jax.checkpoint (jax.remat)https://jax.dev/301/remat.html
Defining new JAX types with hijaxhttps://jax.dev/301/hijax-types.html
JAX 401: kernels and FFIhttps://jax.dev/401/index.html
Pallas: custom kernels in JAXhttps://jax.dev/401/pallas.html
Writing high-performance GPU kernels with CuTe DSL and JAXhttps://jax.dev/401/cute-dsl.html
Foreign function interface (FFI)https://jax.dev/401/ffi.html
JAX 501: systems topicshttps://jax.dev/501/index.html
Introduction to multi-controller JAX (aka multi-process/multi-host JAX)https://jax.dev/501/multiprocess.html
Distributed data loadinghttps://jax.dev/501/data-loading.html
Fault-tolerant distributed JAXhttps://jax.dev/501/fault-tolerance.html
Security considerationshttps://jax.dev/501/security.html
Exporting and serializing staged-out computationshttps://jax.dev/501/export.html
Shape polymorphismhttps://jax.dev/501/shape-polymorphism.html
Persistent compilation cachehttps://jax.dev/501/compilation-cache.html
Transfer guardhttps://jax.dev/501/transfer-guard.html
JAX 601: internalshttps://jax.dev/601/index.html
The jaxpr languagehttps://jax.dev/601/jaxpr.html
Primitiveshttps://jax.dev/601/jax-primitives.html
Autodidax: JAX core from scratchhttps://jax.dev/autodidax.html
Autodidax2, part 1: JAX from scratch, againhttps://jax.dev/autodidax2_part1.html
🔪 JAX - The Sharp Bits 🔪https://jax.dev/notebooks/Common_Gotchas_in_JAX.html
Pallas: a JAX kernel languagehttps://jax.dev/pallas/index.html
Pallas TPUhttps://jax.dev/pallas/tpu/index.html
Quickstart: TPUhttps://jax.dev/pallas/tpu/quickstart.html
Writing TPU kernels with Pallashttps://jax.dev/pallas/tpu/details.html
TPU Pipelininghttps://jax.dev/pallas/tpu/pipelining.html
Matrix Multiplicationhttps://jax.dev/pallas/tpu/matmul.html
Scalar Prefetch and Block-Sparse Computationhttps://jax.dev/pallas/tpu/sparse.html
Distributed Computing in Pallas for TPUshttps://jax.dev/pallas/tpu/distributed.html
Pallas Core-specific Programminghttps://jax.dev/pallas/tpu/core_map.html
SparseCore Kernel Writinghttps://jax.dev/pallas/tpu/sparsecore.html
Fuser API: Fuse operations around kernelshttps://jax.dev/pallas/tpu/fuser.html
Pseudo-Random Number Generationhttps://jax.dev/pallas/tpu/prng.html
TPU Hardware Referencehttps://jax.dev/pallas/tpu/hardware.html
Pallas Mosaic GPUhttps://jax.dev/pallas/gpu/index.html
Quickstart: GPUhttps://jax.dev/pallas/gpu/quickstart.html
Mosaic GPU Pipelininghttps://jax.dev/pallas/gpu/pipelining.html
Writing high-performance matrix multiplication kernels for Blackwellhttps://jax.dev/pallas/gpu/blackwell_matmul.html
Collective matrix multiplicationhttps://jax.dev/pallas/gpu/collective_matmul.html
Writing Mosaic GPU kernels with Pallashttps://jax.dev/pallas/gpu/reference.html
Working with Refshttps://jax.dev/pallas/refs.html
Grids and BlockSpecshttps://jax.dev/pallas/grid_blockspec.html
Software Pipelininghttps://jax.dev/pallas/pipelining.html
API referencehttps://jax.dev/jax.experimental.pallas.html
Pallas TPU (TensorCore)https://jax.dev/jax.experimental.pallas.tpu.html
Pallas TPU (SparseCore)https://jax.dev/jax.experimental.pallas.tpu_sc.html
Pallas MGPUhttps://jax.dev/jax.experimental.pallas.mosaic_gpu.html
Fuserhttps://jax.dev/jax.experimental.pallas.fuser.html
Pallas Design Noteshttps://jax.dev/pallas/design/index.html
Pallas Designhttps://jax.dev/pallas/design/design.html
Pallas Async Operationshttps://jax.dev/pallas/design/async_note.html
Pallas Changeloghttps://jax.dev/pallas/CHANGELOG.html
Additional guideshttps://jax.dev/advanced_guides.html
Colocated Pythonhttps://jax.dev/notebooks/colocated-python.html
The Training Cookbookhttps://jax.dev/the-training-cookbook.html
GPU performance tipshttps://jax.dev/gpu_performance_tips.html
Numerical accuracyhttps://jax.dev/accuracy.html
API Referencehttps://jax.dev/jax.html
jax.numpy modulehttps://jax.dev/jax.numpy.html
jax.numpy.fft.ffthttps://jax.dev/_autosummary/jax.numpy.fft.fft.html
jax.numpy.fft.fft2https://jax.dev/_autosummary/jax.numpy.fft.fft2.html
jax.numpy.fft.fftfreqhttps://jax.dev/_autosummary/jax.numpy.fft.fftfreq.html
jax.numpy.fft.fftnhttps://jax.dev/_autosummary/jax.numpy.fft.fftn.html
jax.numpy.fft.fftshifthttps://jax.dev/_autosummary/jax.numpy.fft.fftshift.html
jax.numpy.fft.hffthttps://jax.dev/_autosummary/jax.numpy.fft.hfft.html
jax.numpy.fft.iffthttps://jax.dev/_autosummary/jax.numpy.fft.ifft.html
jax.numpy.fft.ifft2https://jax.dev/_autosummary/jax.numpy.fft.ifft2.html
jax.numpy.fft.ifftnhttps://jax.dev/_autosummary/jax.numpy.fft.ifftn.html
jax.numpy.fft.ifftshifthttps://jax.dev/_autosummary/jax.numpy.fft.ifftshift.html
jax.numpy.fft.ihffthttps://jax.dev/_autosummary/jax.numpy.fft.ihfft.html
jax.numpy.fft.irffthttps://jax.dev/_autosummary/jax.numpy.fft.irfft.html
jax.numpy.fft.irfft2https://jax.dev/_autosummary/jax.numpy.fft.irfft2.html
jax.numpy.fft.irfftnhttps://jax.dev/_autosummary/jax.numpy.fft.irfftn.html
jax.numpy.fft.rffthttps://jax.dev/_autosummary/jax.numpy.fft.rfft.html
jax.numpy.fft.rfft2https://jax.dev/_autosummary/jax.numpy.fft.rfft2.html
jax.numpy.fft.rfftfreqhttps://jax.dev/_autosummary/jax.numpy.fft.rfftfreq.html
jax.numpy.fft.rfftnhttps://jax.dev/_autosummary/jax.numpy.fft.rfftn.html
jax.scipy modulehttps://jax.dev/jax.scipy.html
jax.scipy.special.i0https://jax.dev/_autosummary/jax.scipy.special.i0.html
jax.scipy.special.i0ehttps://jax.dev/_autosummary/jax.scipy.special.i0e.html
jax.scipy.special.i1https://jax.dev/_autosummary/jax.scipy.special.i1.html
jax.scipy.special.i1ehttps://jax.dev/_autosummary/jax.scipy.special.i1e.html
jax.scipy.special.ivhttps://jax.dev/_autosummary/jax.scipy.special.iv.html
jax.scipy.special.jnhttps://jax.dev/_autosummary/jax.scipy.special.jn.html
jax.scipy.special.knhttps://jax.dev/_autosummary/jax.scipy.special.kn.html
jax.scipy.special.ynhttps://jax.dev/_autosummary/jax.scipy.special.yn.html
jax.scipy.stats.bernoulli.logpmfhttps://jax.dev/_autosummary/jax.scipy.stats.bernoulli.logpmf.html
jax.scipy.stats.bernoulli.pmfhttps://jax.dev/_autosummary/jax.scipy.stats.bernoulli.pmf.html
jax.scipy.stats.bernoulli.cdfhttps://jax.dev/_autosummary/jax.scipy.stats.bernoulli.cdf.html
jax.scipy.stats.bernoulli.ppfhttps://jax.dev/_autosummary/jax.scipy.stats.bernoulli.ppf.html
jax.lax modulehttps://jax.dev/jax.lax.html
jax.random modulehttps://jax.dev/jax.random.html
jax.sharding modulehttps://jax.dev/jax.sharding.html
jax.ad_checkpoint modulehttps://jax.dev/jax.ad_checkpoint.html
jax.debug modulehttps://jax.dev/jax.debug.html
jax.dlpack modulehttps://jax.dev/jax.dlpack.html
jax.distributed modulehttps://jax.dev/jax.distributed.html
jax.dtypes modulehttps://jax.dev/jax.dtypes.html
jax.ffi modulehttps://jax.dev/jax.ffi.html
jax.flatten_util modulehttps://jax.dev/jax.flatten_util.html
jax.image modulehttps://jax.dev/jax.image.html
jax.nn modulehttps://jax.dev/jax.nn.html
jax.nn.initializers modulehttps://jax.dev/jax.nn.initializers.html
jax.ops modulehttps://jax.dev/jax.ops.html
jax.profiler modulehttps://jax.dev/jax.profiler.html
jax.ref modulehttps://jax.dev/jax.ref.html
jax.stages modulehttps://jax.dev/jax.stages.html
jax.test_util modulehttps://jax.dev/jax.test_util.html
jax.tree modulehttps://jax.dev/jax.tree.html
jax.tree_util modulehttps://jax.dev/jax.tree_util.html
jax.typing modulehttps://jax.dev/jax.typing.html
jax.export modulehttps://jax.dev/jax.export.html
jax.extend modulehttps://jax.dev/jax.extend.html
jax.extend.backend modulehttps://jax.dev/jax.extend.backend.html
jax.extend.core modulehttps://jax.dev/jax.extend.core.html
jax.extend.linear_util modulehttps://jax.dev/jax.extend.linear_util.html
jax.extend.lowering modulehttps://jax.dev/jax.extend.lowering.html
jax.extend.mlir modulehttps://jax.dev/jax.extend.mlir.html
jax.extend.pallas modulehttps://jax.dev/jax.extend.pallas.html
jax.extend.random modulehttps://jax.dev/jax.extend.random.html
jax.extend.xla modulehttps://jax.dev/jax.extend.xla.html
jax.example_libraries modulehttps://jax.dev/jax.example_libraries.html
jax.example_libraries.optimizers modulehttps://jax.dev/jax.example_libraries.optimizers.html
jax.example_libraries.stax modulehttps://jax.dev/jax.example_libraries.stax.html
jax.experimental modulehttps://jax.dev/jax.experimental.html
jax.experimental.checkify modulehttps://jax.dev/jax.experimental.checkify.html
jax.experimental.compilation_cache modulehttps://jax.dev/jax.experimental.compilation_cache.html
jax.experimental.custom_partitioning modulehttps://jax.dev/jax.experimental.custom_partitioning.html
jax.experimental.jet modulehttps://jax.dev/jax.experimental.jet.html
jax.experimental.key_reuse modulehttps://jax.dev/jax.experimental.key_reuse.html
jax.experimental.mesh_utils modulehttps://jax.dev/jax.experimental.mesh_utils.html
jax.experimental.multihost_utils modulehttps://jax.dev/jax.experimental.multihost_utils.html
jax.experimental.pallas modulehttps://jax.dev/jax.experimental.pallas.html
Pallas TPU (TensorCore)https://jax.dev/jax.experimental.pallas.tpu.html
Pallas TPU (SparseCore)https://jax.dev/jax.experimental.pallas.tpu_sc.html
Pallas MGPUhttps://jax.dev/jax.experimental.pallas.mosaic_gpu.html
Fuserhttps://jax.dev/jax.experimental.pallas.fuser.html
jax.experimental.random modulehttps://jax.dev/jax.experimental.random.html
jax.experimental.serialize_executable modulehttps://jax.dev/jax.experimental.serialize_executable.html
jax.experimental.sparse modulehttps://jax.dev/jax.experimental.sparse.html
jax.experimental.sparse.BCOOhttps://jax.dev/_autosummary/jax.experimental.sparse.BCOO.html
jax.experimental.sparse.bcoo_broadcast_in_dimhttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_broadcast_in_dim.html
jax.experimental.sparse.bcoo_concatenatehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_concatenate.html
jax.experimental.sparse.bcoo_dot_generalhttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_dot_general.html
jax.experimental.sparse.bcoo_dot_general_sampledhttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_dot_general_sampled.html
jax.experimental.sparse.bcoo_dynamic_slicehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_dynamic_slice.html
jax.experimental.sparse.bcoo_extracthttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_extract.html
jax.experimental.sparse.bcoo_fromdensehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_fromdense.html
jax.experimental.sparse.bcoo_gatherhttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_gather.html
jax.experimental.sparse.bcoo_multiply_densehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_multiply_dense.html
jax.experimental.sparse.bcoo_multiply_sparsehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_multiply_sparse.html
jax.experimental.sparse.bcoo_update_layouthttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_update_layout.html
jax.experimental.sparse.bcoo_reduce_sumhttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_reduce_sum.html
jax.experimental.sparse.bcoo_reshapehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_reshape.html
jax.experimental.sparse.bcoo_slicehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_slice.html
jax.experimental.sparse.bcoo_sort_indiceshttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_sort_indices.html
jax.experimental.sparse.bcoo_squeezehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_squeeze.html
jax.experimental.sparse.bcoo_sum_duplicateshttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_sum_duplicates.html
jax.experimental.sparse.bcoo_todensehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_todense.html
jax.experimental.sparse.bcoo_transposehttps://jax.dev/_autosummary/jax.experimental.sparse.bcoo_transpose.html
jax.experimental.xla_metadata modulehttps://jax.dev/jax.experimental.xla_metadata.html
jax.Array.addressable_shardshttps://jax.dev/_autosummary/jax.Array.addressable_shards.html
jax.Array.allhttps://jax.dev/_autosummary/jax.Array.all.html
jax.Array.anyhttps://jax.dev/_autosummary/jax.Array.any.html
jax.Array.argmaxhttps://jax.dev/_autosummary/jax.Array.argmax.html
jax.Array.argminhttps://jax.dev/_autosummary/jax.Array.argmin.html
jax.Array.argpartitionhttps://jax.dev/_autosummary/jax.Array.argpartition.html
jax.Array.argsorthttps://jax.dev/_autosummary/jax.Array.argsort.html
jax.Array.astypehttps://jax.dev/_autosummary/jax.Array.astype.html
jax.Array.athttps://jax.dev/_autosummary/jax.Array.at.html
jax.Array.byteswaphttps://jax.dev/_autosummary/jax.Array.byteswap.html
jax.Array.choosehttps://jax.dev/_autosummary/jax.Array.choose.html
jax.Array.cliphttps://jax.dev/_autosummary/jax.Array.clip.html
jax.Array.compresshttps://jax.dev/_autosummary/jax.Array.compress.html
jax.Array.committedhttps://jax.dev/_autosummary/jax.Array.committed.html
jax.Array.conjhttps://jax.dev/_autosummary/jax.Array.conj.html
jax.Array.conjugatehttps://jax.dev/_autosummary/jax.Array.conjugate.html
jax.Array.copyhttps://jax.dev/_autosummary/jax.Array.copy.html
jax.Array.copy_to_host_asynchttps://jax.dev/_autosummary/jax.Array.copy_to_host_async.html
jax.Array.cumprodhttps://jax.dev/_autosummary/jax.Array.cumprod.html
jax.Array.cumsumhttps://jax.dev/_autosummary/jax.Array.cumsum.html
jax.Array.devicehttps://jax.dev/_autosummary/jax.Array.device.html
jax.Array.diagonalhttps://jax.dev/_autosummary/jax.Array.diagonal.html
jax.Array.dothttps://jax.dev/_autosummary/jax.Array.dot.html
jax.Array.dtypehttps://jax.dev/_autosummary/jax.Array.dtype.html
jax.Array.flathttps://jax.dev/_autosummary/jax.Array.flat.html
jax.Array.flattenhttps://jax.dev/_autosummary/jax.Array.flatten.html
jax.Array.global_shardshttps://jax.dev/_autosummary/jax.Array.global_shards.html
jax.Array.imaghttps://jax.dev/_autosummary/jax.Array.imag.html
jax.Array.is_fully_addressablehttps://jax.dev/_autosummary/jax.Array.is_fully_addressable.html
jax.Array.is_fully_replicatedhttps://jax.dev/_autosummary/jax.Array.is_fully_replicated.html
jax.Array.itemhttps://jax.dev/_autosummary/jax.Array.item.html
jax.Array.itemsizehttps://jax.dev/_autosummary/jax.Array.itemsize.html
jax.Array.maxhttps://jax.dev/_autosummary/jax.Array.max.html
jax.Array.meanhttps://jax.dev/_autosummary/jax.Array.mean.html
jax.Array.minhttps://jax.dev/_autosummary/jax.Array.min.html
jax.Array.nbyteshttps://jax.dev/_autosummary/jax.Array.nbytes.html
jax.Array.ndimhttps://jax.dev/_autosummary/jax.Array.ndim.html
jax.Array.nonzerohttps://jax.dev/_autosummary/jax.Array.nonzero.html
jax.Array.prodhttps://jax.dev/_autosummary/jax.Array.prod.html
jax.Array.ptphttps://jax.dev/_autosummary/jax.Array.ptp.html
jax.Array.ravelhttps://jax.dev/_autosummary/jax.Array.ravel.html
jax.Array.realhttps://jax.dev/_autosummary/jax.Array.real.html
jax.Array.repeathttps://jax.dev/_autosummary/jax.Array.repeat.html
jax.Array.reshapehttps://jax.dev/_autosummary/jax.Array.reshape.html
jax.Array.roundhttps://jax.dev/_autosummary/jax.Array.round.html
jax.Array.searchsortedhttps://jax.dev/_autosummary/jax.Array.searchsorted.html
jax.Array.shapehttps://jax.dev/_autosummary/jax.Array.shape.html
jax.Array.shardinghttps://jax.dev/_autosummary/jax.Array.sharding.html
jax.Array.sizehttps://jax.dev/_autosummary/jax.Array.size.html
jax.Array.sorthttps://jax.dev/_autosummary/jax.Array.sort.html
jax.Array.squeezehttps://jax.dev/_autosummary/jax.Array.squeeze.html
jax.Array.stdhttps://jax.dev/_autosummary/jax.Array.std.html
jax.Array.sumhttps://jax.dev/_autosummary/jax.Array.sum.html
jax.Array.swapaxeshttps://jax.dev/_autosummary/jax.Array.swapaxes.html
jax.Array.takehttps://jax.dev/_autosummary/jax.Array.take.html
jax.Array.to_devicehttps://jax.dev/_autosummary/jax.Array.to_device.html
jax.Array.tracehttps://jax.dev/_autosummary/jax.Array.trace.html
jax.Array.transposehttps://jax.dev/_autosummary/jax.Array.transpose.html
jax.Array.varhttps://jax.dev/_autosummary/jax.Array.var.html
jax.Array.viewhttps://jax.dev/_autosummary/jax.Array.view.html
jax.Array.Thttps://jax.dev/_autosummary/jax.Array.T.html
jax.Array.mThttps://jax.dev/_autosummary/jax.Array.mT.html
API compatibilityhttps://jax.dev/api_compatibility.html
Python and NumPy version support policyhttps://jax.dev/deprecation.html
Developer noteshttps://jax.dev/contributor_guide.html
Contributing to JAXhttps://jax.dev/contributing.html
Building from sourcehttps://jax.dev/developer.html
Investigating a regressionhttps://jax.dev/investigating_a_regression.html
JAX Enhancement Proposals (JEPs)https://jax.dev/jep/index.html
263: JAX PRNG Designhttps://jax.dev/jep/263-prng.html
2026: Custom JVP/VJP rules for JAX-transformable functionshttps://jax.dev/jep/2026-custom-derivatives.html
4008: Custom VJP and `nondiff_argnums` updatehttps://jax.dev/jep/4008-custom-vjp-update.html
4410: Omnistaginghttps://jax.dev/jep/4410-omnistaging.html
9263: Typed keys & pluggable RNGshttps://jax.dev/jep/9263-typed-keys.html
9407: Design of Type Promotion Semantics for JAXhttps://jax.dev/jep/9407-type-promotion.html
9419: Jax and Jaxlib versioninghttps://jax.dev/jep/9419-jax-versioning.html
10657: Sequencing side-effects in JAXhttps://jax.dev/jep/10657-sequencing-effects.html
11830: `jax.remat` / `jax.checkpoint` new implementationhttps://jax.dev/jep/11830-new-remat-checkpoint.html
12049: Type Annotation Roadmap for JAXhttps://jax.dev/jep/12049-type-annotations.html
14273: `shard_map` (`shmap`) for simple per-device codehttps://jax.dev/jep/14273-shard-map.html
15856: `jax.extend`, an extensions modulehttps://jax.dev/jep/15856-jex.html
17111: Efficient transposition of `shard_map` (and other maps)https://jax.dev/jep/17111-shmap-transpose.html
18137: Scope of JAX NumPy & SciPy Wrappershttps://jax.dev/jep/18137-numpy-scipy-scope.html
25516: Effort-based versioninghttps://jax.dev/jep/25516-effver.html
28661: Supporting the `__jax_array__` protocolhttps://jax.dev/jep/28661-jax-array-protocol.html
28845: Stateful Randomness in JAXhttps://jax.dev/jep/28845-stateful-rng.html
JAX Internal Implementation Noteshttps://jax.dev/internals/index.html
Handling of closed-over constantshttps://jax.dev/internals/constants.html
Extension guideshttps://jax.dev/extensions.html
Writing custom Jaxpr interpreters in JAXhttps://jax.dev/notebooks/Writing_custom_interpreters_in_Jax.html
jax.extend modulehttps://jax.dev/jax.extend.html
jax.extend.backend modulehttps://jax.dev/jax.extend.backend.html
jax.extend.core modulehttps://jax.dev/jax.extend.core.html
jax.extend.linear_util modulehttps://jax.dev/jax.extend.linear_util.html
jax.extend.lowering modulehttps://jax.dev/jax.extend.lowering.html
jax.extend.mlir modulehttps://jax.dev/jax.extend.mlir.html
jax.extend.pallas modulehttps://jax.dev/jax.extend.pallas.html
jax.extend.random modulehttps://jax.dev/jax.extend.random.html
jax.extend.xla modulehttps://jax.dev/jax.extend.xla.html
Building on JAXhttps://jax.dev/building_on_jax.html
About the projecthttps://jax.dev/about.html
Frequently asked questions (FAQ)https://jax.dev/faq.html
Change loghttps://jax.dev/changelog.html
Glossary of termshttps://jax.dev/glossary.html
Configuration Optionshttps://jax.dev/config_options.html
https://github.com/jax-ml/jax
.rst https://jax.dev/_sources/index.rst
Documentationhttps://jax.dev#documentation
Ecosystemhttps://jax.dev#ecosystem
#https://jax.dev#jax-high-performance-array-computing
Installationhttps://jax.dev/installation.html#installation
JAX 101: expressing computationshttps://jax.dev/101/index.html#jax-101
API Referencehttps://jax.dev/jax.html
JAX AI Stackhttps://docs.jaxstack.ai/en/latest/getting_started.html
#https://jax.dev#documentation
JAX 101https://jax.dev/101/index.html
jax.numpyhttps://jax.dev/jax.numpy.html#module-jax.numpy
jax.grad()https://jax.dev/_autosummary/jax.grad.html#jax.grad
jax.vmap()https://jax.dev/_autosummary/jax.vmap.html#jax.vmap
JAX 201https://jax.dev/201/index.html
jax.jit()https://jax.dev/_autosummary/jax.jit.html#jax.jit
JAX 301https://jax.dev/301/index.html
JAX 401https://jax.dev/401/index.html
JAX 501https://jax.dev/501/index.html
JAX 601https://jax.dev/601/index.html
Newly documentedhttps://jax.dev/whats-new.html
#https://jax.dev#ecosystem
Flaxhttps://flax.readthedocs.io/
Equinoxhttps://docs.kidger.site/equinox/
Kerashttps://keras.io/
Optaxhttps://optax.readthedocs.io/
Optimistixhttps://github.com/patrick-kidger/optimistix
Lineaxhttps://github.com/patrick-kidger/lineax
Diffraxhttps://docs.kidger.site/diffrax/
Grainhttps://github.com/google/grain
TensorFlow Datasetshttps://www.tensorflow.org/datasets
Hugging Face Datasetshttps://huggingface.co/docs/datasets/
Orbaxhttps://orbax.readthedocs.io/
Chexhttps://chex.readthedocs.io/
Blackjaxhttps://blackjax-devs.github.io/blackjax/
Numpyrohttps://num.pyro.ai/en/latest/index.html
PyMChttps://www.pymc.io/
TensorFlow Probabilityhttps://www.tensorflow.org/probability
Distraxhttps://github.com/google-deepmind/distrax
JAX MDhttps://jax-md.readthedocs.io/
Braxhttps://github.com/google/brax/
MaxTexthttps://github.com/google/maxtext/
AXLearnhttps://github.com/apple/axlearn
Levanterhttps://github.com/stanford-crfm/levanter
EasyLMhttps://github.com/young-geng/EasyLM
Marinhttps://github.com/marin-community/marin
Awesome JAXhttps://github.com/n2cholas/awesome-jax
next Installation https://jax.dev/installation.html
Documentationhttps://jax.dev#documentation
Ecosystemhttps://jax.dev#ecosystem

Viewport: width=device-width, initial-scale=1

URLs of crawlers that visited me.