Merge pull request #33424 from jakevdp:error-doc

PiperOrigin-RevId: 834767465

Author: Google-ML-Automation

docs/errors.rst

@@ -7,9 +7,11 @@ This page lists a few of the errors you might encounter when using JAX,
 along with representative examples of how one might fix them.
 
 .. currentmodule:: jax.errors
+.. autoclass:: JaxRuntimeError
+.. autoclass:: JAXTypeError
+.. autoclass:: JAXIndexError
 .. autoclass:: ConcretizationTypeError
 .. autoclass:: KeyReuseError
-.. autoclass:: JaxRuntimeError
 .. autoclass:: NonConcreteBooleanIndexError
 .. autoclass:: TracerArrayConversionError
 .. autoclass:: TracerBoolConversionError

jax/_src/errors.py

@@ -35,12 +35,12 @@ def __init__(self, message: str):
 
 @export
 class JAXTypeError(_JAXErrorMixin, TypeError):
-  pass
+  """JAX-specific :class:`TypeError`"""
 
 
 @export
 class JAXIndexError(_JAXErrorMixin, IndexError):
-  pass
+  """JAX-specific :class:`IndexError`"""
 
 
 @export

tests/documentation_coverage_test.py

@@ -58,7 +58,6 @@ def jax_docs_dir() -> str:
   'jax.debug': ['DebugEffect', 'log'],
   'jax.distributed': ['is_initialized'],
   'jax.dtypes': ['extended', 'finfo', 'iinfo'],
-  'jax.errors': ['JAXIndexError', 'JAXTypeError'],
   'jax.ffi': ['build_ffi_lowering_function', 'include_dir', 'register_ffi_target_as_batch_partitionable', 'register_ffi_type_id'],
   'jax.lax': ['unreduced_psum', 'dce_sink', 'conv_transpose_shape_tuple', 'reduce_window_shape_tuple', 'preduced', 'conv_general_permutations', 'conv_general_shape_tuple', 'pbroadcast', 'padtype_to_pads', 'conv_shape_tuple', 'unreduced_psum_scatter', 'create_token', 'dtype', 'shape_as_value', 'all_gather_reduced', 'pvary', *(name for name in dir(jax.lax) if name.endswith('_p'))],
   'jax.lax.linalg': [api for api in dir(jax.lax.linalg) if api.endswith('_p')],