Skip to main content

Ctrl+K

Getting Started

Installing JAX
JAX Quickstart
How to Think in JAX
🔪 JAX - The Sharp Bits 🔪
JAX Frequently Asked Questions (FAQ)

Tutorial: JAX 101

Further Resources

User Guides
Advanced Tutorials
Developer Documentation
Building on JAX
Notes
Public API: jax package

Change log
JAX Glossary of Terms

.rst

jax.lax.after_all

Contents

after_all()

jax.lax.after_all#

jax.lax.after_all(*operands)[source]#

Merges one or more XLA token values. Experimental.

Wraps the XLA AfterAll operator.

previous

jax.lax.add

next

jax.lax.approx_max_k

Contents

after_all()

By The JAX authors

© Copyright 2024, The JAX Authors. NumPy and SciPy documentation are copyright the respective authors..