Skip to main content

Ctrl+K

Getting started

Installation
Quickstart: How to think in JAX

🔪 JAX - The Sharp Bits 🔪
JAX 101

Resources, guides, and references

Key concepts
Resources and Advanced Guides
API Reference
Developer notes
Extension guides
Notes
Pallas: a JAX kernel language
About the project

Frequently asked questions (FAQ)
Change log
Glossary of terms

Configuration Options

API Reference
jax.extend module
jax.extend.core module
jax.extend.core.unsafe_get_axis_names_DO_NOT_USE

.rst

jax.extend.core.unsafe_get_axis_names_DO_NOT_USE

Contents

unsafe_get_axis_names_DO_NOT_USE()

jax.extend.core.unsafe_get_axis_names_DO_NOT_USE#

jax.extend.core.unsafe_get_axis_names_DO_NOT_USE()[source]#

Return type:: list[Any]

previous

jax.extend.core.unsafe_am_i_under_a_vmap_DO_NOT_USE

next

jax.extend.core.valid_jaxtype

Contents

unsafe_get_axis_names_DO_NOT_USE()

By The JAX authors

© Copyright 2024, The JAX Authors.