 Define a class that contains all the data needed to calculate
        a detection efficiency (DE) model
        The contents of this class can largely be treated as a black box
        as they are filled in by the init function and are very dependent
        on the data tables that are used to build the per-target
        DE model.  More information on this model is available in
        Burke, C.J. & Catanzarite, J. 2017, "Planet Detection Metrics:
           Per-Target Detection Contours for Data Release 25", KSCI-19111-001

        INIT INPUT:
        filePath - [str] path to the directory containing fits files

        CLASS VARIABLE CONTENTS:
        lowPers, hghPers, midpers, mesValues, fitOrder, fitCoeffs
           These get set by the contents of detectEffData_alpha12_02272017.h5
           These are associated with the period dependent coeffcients for the DE model
        baseDetFunc
           This is a function that interpolates the Rstar dependent 'base' DE model
           as stored in detectEffData_alpha_base_02272017.txt
        rStarRange - Valid range over Rstar for Rstar 'base' DE model
        cSLSSDetFunc
           This is a function that interpolates the CDPP slope plane DE corrections
           as stored in detectEffData_alpha12_SlopeLongShort_02272017.txt
        cdppSlopeLongRange/cdppSlopeShortRange
           Valid range of the short and long cdpp slope parameters.
        detEffValues - Tabulated DE values for target dependent DE model
        detEffFunctions - Interpolation function for non-smeared DE model
        mesSmearA, mesSmearB - Beta distribution function parameters for MES smearing
        detEffSmearFunctions - Final DE model after MES Smearing applied

        CLASS FUNCTION CONTENTS:
        __init__ - initialization and read in all necessary tables and make
                   DE model interpolating functions
        generate_detection_eff_grid - After instantiation call this function
                to tailor the DE model for the targets parameters
        mes_smearing_parameters - Get the beta distribution parameters
                for the mes smearing model
        final_detEffs - This is the function the end user will actually
                use in order to get a DE model for any arbitary MES and period


Define a class that contains all the data needed to calculate
       a single target pipeline completeness grid using
       kepler_single_comp()
       CONTENTS:
       id - [int] Target identifier recommend KIC
       period_want - [day] list of orbital periods
       rp_want - [Rearth] list of planet radii
       rstar - [Rsun] star radius
       logg - [cgs] star surface gravity
       teff - [K] stellar effective temperature
       ecc - [0.0 - 1.0] orbital eccentricity
       dataspan - [day] scalar observing baseline duration
       dutycycle - [0.0 -1.0] scalar valid data fraction over dataspan
       pulsedurations - [hr] list of transit durations searched
       cdppSlopeLong - [float] rmsCDPP Slope for long durations
       cdppSlopeShort - [float] rmsCDPP Slope for short durations
       planet_detection_metric_path - [string] directory path
                                        of fits files for the
                                        planet detection metrics


 Arrays allocated on shared memory.

        The array is stored in an anonymous memory map that is shared between child-processes.


    Profile class

    Parameters
    ----------
    counts : array-like
        the counts in each phase bin of Profile

    cycles : int
        the period cycles of input Profile (default is 1).
        If cycles=2, the phase of profile would be ``np.linspace(0, 2, size_of_Profile+1)[:-1]``

    error : array-like
        the error of each phase bin, if not given the error will be the
        poisson error of counts (sqruare root of counts)

    Attributes
    ----------
    counts : array-like
        The counts in each phase bin of Profile
    phase : array-like
        The midpoints of phase bins
    phase_off : list
        The list of phase off interval, the two value are the left
        and right phase bin of off pulse phases.
        left_edge = phase_off[0]
        right_edge = phase_ff[1]
    ref_time : float, default None
        The reference time (in second) used to fold the profile


    The parameter in Timing Model. One can retrieve the value and/or error (if valid)
    of the parameter.

    Examples
    --------
    >>> f0 = TimingParameter(value=27, error=1e-3, name='F0')
    >>> print(f'{f0.name} = {f0.value} +- {f0.error}')
    >>> F0 = 27 +- 0.001


    Class of Pulsar timing model

    Examples
    --------
    >>> eph = TimingModel()
    >>> eph.readpar("test.par")
    >>> print("F0 = ", eph.F0.value, eph.f0.value)
    >>> F0 =  29.636679699921209437
    >>> print("F0 error = ", eph.F0.error)
    >>> F0 error =  1.7247236495e-09
    >>> print("PEPOCH = ", eph.PEPOCH.value)
    >>> PEPOCH =  58066.18087539147382
    >>> print("PEPOCH error = ", eph.PEPOCH.error)
    >>> PEPOCH error =  None


Abstract class for a z-tomography of a single datacube


Abstract class for a z-tomography of 2 datacubes and their residual


z-Tomography of a datacube representing photoionization rates


z-Tomography of a datacube representing ionization/neutral fractions


z-Tomography of 2 photoionization rate datacubes and their residual


z-Tomography of 2 ionization/neutral fraction datacubes and their residual


A point source emitting a Black-body spectrum


Flexible test generator class.

    Example usage:

     x = TestGenerator()
     x.go()

    This runs genetIC twice, once with actual zoom regions and another time with an equivalent uniform-res grid
    then plots a four-panel plot to show the result and differences.

    To customise, you can either derive a class and change its settings or manually manipulate the object, e.g.
    by setting x.zoom_factor = [2,2]; x.zoom_ncells = [64, 64] you will get a double zoom.

    For all parameters that can be varied look at the top of the TestGenerator class.


Calculate the low-res/high-res split by making a full box at the high resolution,
    then downgrading the resolution of the low-res region


The original implementation of fast-filter zooms, which corresponds more closely to the genetIC code but
    shifts various factors around and therefore looks quite different from the FilteredZoomConstrained implementation
    above.


The reference implementation for fast-filter zooms, which should match exactly the description in the notes


A version of FilteredZoomConstrained that works in the delta_Z basis instead of the noise_Z basis


This is a trial implementation to see if a hybrid approach can improve accuracy further (it doesn't seem to help)


A base class for those zoom implementations which adopt the covector/vector geometry described in the notes.


Test the `generate.py` script


Tests for the metadata utils module used to convert between parameter definitions


Tests for the module used to create PSFs.


Tests for the selection module used to accept or reject sampled parameters


A suite of tests verifying that the input PDFs and the sample distributions
    match.


A suite of tests alerting us for breakge, e.g. errors in
    instantiation of classes or execution of scripts, for DiagonalBNNPrior


A suite of tests alerting us for breakge, e.g. errors in
    instantiation of classes or execution of scripts, for CovBNNPrior


A suite of tests alerting us for breakge, e.g. errors in
    instantiation of classes or execution of scripts, for EmpiricalBNNPrior


Nested dictionary representing the configuration for Baobab data generation


A combination of sky, readout, and Poisson flux noise to be added to the image

    Note
    ----
    This is a torch wrapper around the functionality provided by the `SingleBand` class in lenstronomy.


A combination of sky, readout, and Poisson flux noise to be added to the image

    Note
    ----
    This is a wrapper around the functionality provided by the `SingleBand` class in lenstronomy.


Abstract base class for a cosmology-aware BNN prior


BNN prior with independent parameters

    Note
    ----
    This BNNPrior is cosmology-agnostic. For a version that's useful for H0 inference, see `DiagonalCosmoBNNPrior`.


BNN prior with independent parameters

    Note
    ----
    This BNNPrior is cosmology-aware. For a version that's not tailored to H0 inference, see `DiagonalBNNPrior`.


Represents the Faber-Jackson (FJ) relation between velocity dispersion and luminosity
	of elliptical galaxies.

	FJ is a projection of the Fundamental Plane (FP) relation.


Represents the Fundamental Plane (FP) relation between the velocity dispersion,
	luminosity, and effective radius for elliptical galaxies

	Luminosity is expressed as apparent magnitude in this form.


Represents bivariate relations (projections) within the Fundamental Mass Hyperplane (FMHP) relation
	between the stellar mass, stellar mass density, effective radius, and velocity dispersion of massive ETGs.

	Only the relation between the power-law mass slope (gamma) and effective radius is currently supported.


Represents various scaling relations that the axis ratio can follow with
	quantities like velocity dispersion, when its PDF is assumed to be a Rayleigh distribution

	Only the relation with velocity dispersion is currently supported.


Redshift-binned AGN luminosity function parameterized as a double power-law


BNN prior that encodes physical correlations between parameters


Abstract base class equipped with PDF evaluation and sampling utility functions for various lens/source macromodels


BNN prior with marginally covariant parameters

    Note
    ----
    This BNNPrior is cosmology-agnostic. For a version that's useful for H0 inference, see `CovCosmoBNNPrior`.


r"""Applies K correction, i.e. "corrects" for the redshift and bandpass throughput in the conversion between apparent and absolute magnitudes

    Note
    ----
    The following is summarized from [1]_. Consider a source observed to have apparent magnitude :math:`m_R` when observed through photometric bandpass :math:`R`, for which we want to know the absolute magnitude :math:`M_Q` in emitted-frame bandpass :math:`Q`. Denote the distance modulus as :math:`DM`. The K correction term :math:`K_{QR}` for this source is defined as

    .. math::

        m_R = M_Q + DM + K_{QR}

    We can compute :math:`K_{QR}` via

    .. math::

        K_{QR} = -2.5 \log_{10} \left[ \frac{1}{1 + z} \frac{\int d\lambda_0 \lambda_0 f_\lambda(\lambda_0) R(\lambda_0) \int d\lambda_e \lambda_e g_\lambda^Q(\lambda_e) Q(\lambda_e) }{ \int d\lambda_0 \lambda_0 g_\lambda^R(\lambda_0) R(\lambda_0) \int d\lambda_e f_\lambda\left((1 + z)\lambda_e \right) Q(\lambda_e)} \right]

    where :math:`f_\lambda` is the spectral density of flux and the superscript denotes

    References
    ----------
    .. [1] Hogg, David W., et al. "The K correction." arXiv preprint astro-ph/0210394 (2002).


General class for chromatic objects, or objects carrying wavelength-dependent properties


Represents the instrument bandpass, characterized by its throughput


Represents the spectral energy distribution (SED) of a source


Dev-mode class, more flexible than Imager.

    Note
    ----
    Accompanying `generate` script doesn't exist yet.


Deterministic utility class for imaging the objects on a pixel grid

        Attributes
        ----------
        bnn_omega : dict
            copy of `cfg.bnn_omega`
        components : list
            list of components, e.g. `lens_mass`


Selections applied to the sampled set of parameters


Tests for LAX-backed scipy.interpolate implementations


Shared tests between the Python and the C++ jax,jit implementations.

  Because the Python implementation supports more features, we need to have the
  Python tests that extend the C++ tests (and not the other way around).


Numerical tests for LAX operations.


Tests of transformations on randomly generated functions.


Tests for LAX-backed scipy.spatial implementations


Tests a function with shape polymorphism.

  Exports `fun` with shape polymorphism, then checks that the JAX native and
  the exported function produce the same results.


This test runs for all _POLY_SHAPE_PRIMITIVE_HARNESSES.


Tests for LAX-backed scipy.stats implementations


Tests for LAX-backed Numpy implementation.


Tests for LAX-backed Numpy operators.


Class that provides a __jax_array__ method.


Duck-typed array that provides a __jax_array__ method which fails.


Shortcut for specifying ShapeDtypeStruct.


Tests for lax.linalg primitives.


Stores a value but pretends to be equal to every other black box.


Simple smoke-tests for tree_util aliases under jax.tree


Mock TPU device for testing.


Tests for Numpy indexing translation rules.


Smoke test for the array API.


Test that our array API implementations raise errors where required


Use to inspect tracers.

  We can `append` tracers from tracing contexts to this object. We collect
  the tracer, along with the error message we get when we try to concretize
  it. This is meant to simulate errors like concretization or leaking.


Tests for LAX-backed Scipy implementation.


Tests for LAX-backed scipy.fft implementations


An in-memory cache for testing purposes.


Tests jit targeting to different backends.


Test interactions among the custom_{vmap,jvp,vjp,transpose,*} APIs


Tests of common functionality that should be run with all PRNG impls.


  Tests of distribution statistics that need only be run with the default PRNG.

  We limit this to the default PRNG to avoid repeated execution of very costly
  tests. So long as the input bits are valid (as tested in BasicRandomTest) then
  the distribution logic tested here will apply correctly.


Tests for LAX-backed Numpy reduction operations.


Tests for loading/storing from bool memrefs on TPUs.

  We specifically test bools because they have special handling.
  Bools are stored as integers inside of memrefs, and we typecast to/from
  bools automatically on load.


Tests for reported bugs. Only pass in interpret mode unless fixed.


Identity mask that is guaranteed to trigger grid shrinking.


Check the construction of MaskInfo from Mask.


These are unit tests for the indexer logic, not using pallas_call.


Wraps the I/O callback `store` into a callback that counts the number of calls to `store`.


Records grid points in the order in which they are procsessed.


A type that represents a 8-bit signed integer.

  This is a workaround to bypass the fact that we don't have a proper 8-bit
  integer type class available in MLIR, and can't instantiate types without a
  MLIR context.


Defines a Tile transform in a TestCaseInput.

  Note that we cannot simply alias mgpu_dialect.TileTransformAttr.get, because
  we do not have an MLIR context at the point we define the TestCaseInput.


Defines a Transpose transform in a TestCaseInput.

  Note that we cannot simply alias mgpu_dialect.TransposeTransformAttr.get,
  because we do not have an MLIR context at the point we define the
  TestCaseInput.


Defines a Swizzle transform in a TestCaseInput.

  Note that we cannot simply alias mgpu_dialect.SwizzleTransformAttr.get,
  because we do not have an MLIR context at the point we define the
  TestCaseInput.


Device tests with lowering from the MLIR dialect and layout inference.


Represents the result of executing a subprocess command.


Manages execution of subprocess commands with reusable environment and logging.


A noop variant of the ``legacy`` directive from SciPy.


This class makes 'bdist_wheel' include an ABI tag on the wheel.


Raised when the GPU library is not linked.


An extension to the automatically generated TraceOp bindings.


An extension to the automatically generated RegionOp bindings.


Marks the boundary between two joined (nested) pytrees.


Customizes the behavior of cloudpickle.


Redefine the internal error state based on the mesh in the context.

  When using JAX in multi-device environments in explicit mode, error tracking
  needs to be properly aligned with the device mesh. This context manager
  ensures that the internal error state is correctly initialized based on the
  current mesh configuration.

  This context manager should be used when starting a multi-device computation,
  or when switching between different device meshes.

  On entering the context, it initializes a new error state based on the mesh in
  the context. On exiting the context, it restores the previous error state.


A class to store error information for AOT compilation.

  This class is used internally by the wrapper functions `wrap_for_export` and
  `unwrap_from_import` to encapsulate error-related data within an exported
  function.

  Attributes:
    error_code (jax.Array): A JAX array representing the final error state of
      the function to be exported. This value is local to the wrapper function.
    error_list (list[tuple[str, str]]): A list of `(error_message, traceback)`
      pairs containing error messages and corresponding stack traces. This error
      list is local to the wrapper function, and does not contain pairs of error
      information from other functions.


Customize the vmap behavior of a JAX-transformable function.

  This decorator is used to customize the behavior of a JAX function under the
  :func:`jax.vmap` transformation. A ``custom_vmap``-decorated function will
  mostly (see below for caveats) have the same behavior as the underlying
  function, except when batched using :py:func:`jax.vmap`. When batched, the
  rule defined using :py:func:`~jax.custom_batching.custom_vmap.def_vmap` will
  be used.

  For example:

    >>> @jax.custom_batching.custom_vmap
    ... def f(x, y):
    ...   return x + y
    ...
    >>> @f.def_vmap
    ... def f_vmap_rule(axis_size, in_batched, xs, ys):
    ...   assert all(in_batched)
    ...   assert xs.shape[0] == axis_size
    ...   assert ys.shape[0] == axis_size
    ...   out_batched = True
    ...   return xs * ys, out_batched
    ...
    >>> xs = jnp.arange(3)
    >>> ys = jnp.arange(1, 4)
    >>> jax.vmap(f)(xs, ys)  # prints xs * ys instead of xs + ys
    Array([0, 2, 6], dtype=int32)

  Of note, ``custom_vmap`` functions do not support reverse-mode autodiff. To
  customize both vmap and reverse-mode autodiff, combine ``custom_vmap`` with
  :py:class:`jax.custom_vjp`. For example:

    >>> @jax.custom_vjp
    ... @jax.custom_batching.custom_vmap
    ... def f(x, y):
    ...   return jnp.sin(x) * y
    ...
    >>> @f.def_vmap
    ... def f_vmap_rule(axis_size, in_batched, xs, ys):
    ...   return jnp.cos(xs) * ys, True
    ...
    >>> def f_fwd(x, y):
    ...   return f(x, y), (jnp.cos(x), jnp.sin(x), y)
    ...
    >>> def f_bwd(res, g):
    ...   cos_x, sin_x, y = res
    ...   return (cos_x * g * y, sin_x * g)
    ...
    >>> f.defvjp(f_fwd, f_bwd)
    >>> jax.vmap(f)(jnp.zeros(3), jnp.ones(3))
    Array([1., 1., 1.], dtype=float32)
    >>> jax.grad(f)(jnp.zeros(()), jnp.ones(()))
    Array(1., dtype=float32)

  Note that the :py:class:`jax.custom_vjp` must be on the outside, wrapping the
  ``custom_vmap``-decorated function.


Specifies PRNG key shape and operations.

  A PRNG implementation is determined by a key type ``K`` and a
  collection of functions that operate on such keys. The key type
  ``K`` is an array type with element type uint32 and shape specified
  by ``key_shape``. The type signature of each operations is::

    seed :: int[] -> K
    fold_in :: K -> int[] -> K
    split[shape] :: K -> K[*shape]
    random_bits[shape, bit_width] :: K -> uint<bit_width>[*shape]

  A PRNG implementation is adapted to an array-like object of keys
  ``K`` by the ``PRNGKeyArray`` class, which should be created via the
  ``random_seed`` function.


An array of PRNG keys backed by an RNG implementation.

  This class lifts the definition of a PRNG, provided in the form of a
  ``PRNGImpl``, into an array-like pytree class. Instances of this
  class behave like an array whose base elements are keys, hiding the
  fact that keys are typically arrays (of ``uint32`` dtype) themselves.

  PRNGKeyArrays are also restricted relative to JAX arrays in that
  they do not expose arithmetic operations. They instead expose
  wrapper methods around the PRNG implementation functions (``split``,
  ``random_bits``, ``fold_in``).


Things that we know about a jit instance before it is called.

  In other words, this structure contains arguments to jit()/pjit(),
  preprocessed and validated.


Mutable value object for _infer_params_cached.


Represents an unserialized backend config for custom calls.


r"""A :class:`NamedSharding` expresses sharding using named axes.

  A :class:`NamedSharding` is a pair of a :class:`Mesh` of devices and
  :class:`PartitionSpec` which describes how to shard an array across that
  mesh.

  A :class:`Mesh` is a multidimensional NumPy array of JAX devices,
  where each axis of the mesh has a name, e.g. ``'x'`` or ``'y'``.

  A :class:`PartitionSpec` is a tuple, whose elements can be a ``None``,
  a mesh axis, or a tuple of mesh axes. Each element describes how an input
  dimension is partitioned across zero or more mesh dimensions. For example,
  ``PartitionSpec('x', 'y')`` says that the first dimension of data
  is sharded across ``x`` axis of the mesh, and the second dimension is sharded
  across ``y`` axis of the mesh.

  The `Distributed arrays and automatic parallelization`_
  and `Explicit Sharding`_ tutorials have more details and diagrams that
  explain how :class:`Mesh` and :class:`PartitionSpec` are used.

  Args:
    mesh: A :class:`jax.sharding.Mesh` object.
    spec: A :class:`jax.sharding.PartitionSpec` object.

  Examples:

    >>> from jax.sharding import Mesh
    >>> from jax.sharding import PartitionSpec as P
    >>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))
    >>> spec = P('x', 'y')
    >>> named_sharding = jax.sharding.NamedSharding(mesh, spec)

  .. _Distributed arrays and automatic parallelization: https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
  .. _Explicit Sharding:  https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html


Storage for a value, with checks for overwriting or reading empty store.


Represents a function `f` to which `transforms` are to be applied.

  Args:
    f: the function to be transformed.
    f_transformed: transformed function.
    transforms: a tuple of `(gen, gen_static_args)` tuples representing
      transformations to apply to `f.` Here `gen` is a generator function and
      `gen_static_args` is a tuple of static arguments for the generator. See
      description at the start of this module for the expected behavior of the
      generator.
    stores: a list of out_store for the auxiliary output of the `transforms`.
    params: a tuple of `(name, param)` tuples representing extra parameters to
      pass as keyword arguments to `f`, along with the transformed keyword
      arguments.
    in_type: optional input type
    debug_info: debugging info about the function being wrapped.


Debugging info about a func, its arguments, and results.


Stores an unchanging value. Checks empty reads and unequal overwrites.


A mixin with methods allowing to skip arch specific tests.


Base class for JAX tests including numerical checks and boilerplate.


A class that unifies lists of supported dtypes.

  These could be module-level constants, but device_under_test() is not always
  known at import time, so we need to define these lists lazily.


Same as numpy.vectorize but using mpmath backend for function evaluation.


Namespace of universal functions on numpy arrays that use mpmath
  backend for evaluation and return numpy arrays as outputs.


Mixin for JAX-specific errors


  This error occurs when a JAX Tracer object is used in a context where a
  concrete value is required (see :ref:`faq-different-kinds-of-jax-values`
  for more on what a Tracer is). In some situations, it can be easily fixed by
  marking problematic values as static; in others, it may indicate that your
  program is doing operations that are not directly supported by JAX's JIT
  compilation model.

  Examples:

  Traced value where static value is expected
    One common cause of this error is using a traced value where a static value
    is required. For example:

      >>> from functools import partial
      >>> from jax import jit
      >>> import jax.numpy as jnp
      >>> @jit
      ... def func(x, axis):
      ...   return x.min(axis)

      >>> func(jnp.arange(4), 0)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      ConcretizationTypeError: Abstract tracer value encountered where concrete
      value is expected: axis argument to jnp.min().

    This can often be fixed by marking the problematic argument as static::

        >>> @partial(jit, static_argnums=1)
        ... def func(x, axis):
        ...   return x.min(axis)

        >>> func(jnp.arange(4), 0)
        Array(0, dtype=int32)

  Shape depends on Traced Value
    Such an error may also arise when a shape in your JIT-compiled computation
    depends on the values within a traced quantity. For example::

      >>> @jit
      ... def func(x):
      ...     return jnp.where(x < 0)

      >>> func(jnp.arange(4))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected:
      The error arose in jnp.nonzero.

    This is an example of an operation that is incompatible with JAX's JIT
    compilation model, which requires array sizes to be known at compile-time.
    Here the size of the returned array depends on the contents of `x`, and such
    code cannot be JIT compiled.

    In many cases it is possible to work around this by modifying the logic used
    in the function; for example here is code with a similar issue::

      >>> @jit
      ... def func(x):
      ...     indices = jnp.where(x > 1)
      ...     return x[indices].sum()

      >>> func(jnp.arange(4))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      ConcretizationTypeError: Abstract tracer value encountered where concrete
      value is expected: The error arose in jnp.nonzero.

    And here is how you might express the same operation in a way that avoids
    creation of a dynamically-sized index array::

      >>> @jit
      ... def func(x):
      ...   return jnp.where(x > 1, x, 0).sum()

      >>> func(jnp.arange(4))
      Array(5, dtype=int32)

  To understand more subtleties having to do with tracers vs. regular values,
  and concrete vs. abstract values, you may want to read
  :ref:`faq-different-kinds-of-jax-values`.


  This error occurs when a program attempts to use non-concrete boolean indices
  in a traced indexing operation. Under JIT compilation, JAX arrays must have
  static shapes (i.e. shapes that are known at compile-time) and so boolean
  masks must be used carefully. Some logic implemented via boolean masking is
  simply not possible in a :func:`jax.jit` function; in other cases, the logic
  can be re-expressed in a JIT-compatible way, often using the three-argument
  version of :func:`~jax.numpy.where`.

  Following are a few examples of when this error might arise.

  Constructing arrays via boolean masking
    This most commonly arises when attempting to create an array via a boolean
    mask within a JIT context. For example::

      >>> import jax
      >>> import jax.numpy as jnp

      >>> @jax.jit
      ... def positive_values(x):
      ...   return x[x > 0]

      >>> positive_values(jnp.arange(-5, 5))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])

    This function is attempting to return only the positive values in the input
    array; the size of this returned array cannot be determined at compile-time
    unless `x` is marked as static, and so operations like this cannot be
    performed under JIT compilation.

  Reexpressible Boolean Logic
    Although creating dynamically sized arrays is not supported directly, in
    many cases it is possible to re-express the logic of the computation in
    terms of a JIT-compatible operation. For example, here is another function
    that fails under JIT for the same reason::

      >>> @jax.jit
      ... def sum_of_positive(x):
      ...   return x[x > 0].sum()

      >>> sum_of_positive(jnp.arange(-5, 5))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])

    In this case, however, the problematic array is only an intermediate value,
    and we can instead express the same logic in terms of the JIT-compatible
    three-argument version of :func:`jax.numpy.where`::

      >>> @jax.jit
      ... def sum_of_positive(x):
      ...   return jnp.where(x > 0, x, 0).sum()

      >>> sum_of_positive(jnp.arange(-5, 5))
      Array(10, dtype=int32)

    This pattern of replacing boolean masking with three-argument
    :func:`~jax.numpy.where` is a common solution to this sort of problem.

  Boolean indexing into JAX arrays
    The other situation where this error often arises is when using boolean
    indices, such as with :code:`.at[...].set(...)`. Here is a simple example::

      >>> @jax.jit
      ... def manual_clip(x):
      ...   return x.at[x < 0].set(0)

      >>> manual_clip(jnp.arange(-2, 2))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[4])

    This function is attempting to set values smaller than zero to a scalar fill
    value. As above, this can be addressed by re-expressing the logic in terms
    of :func:`~jax.numpy.where`::

      >>> @jax.jit
      ... def manual_clip(x):
      ...   return jnp.where(x < 0, 0, x)

      >>> manual_clip(jnp.arange(-2, 2))
      Array([0, 0, 0, 1], dtype=int32)


  This error occurs when a program attempts to convert a JAX Tracer object into
  a standard NumPy array (see :ref:`faq-different-kinds-of-jax-values` for more
  on what a Tracer is). It typically occurs in one of a few situations.

  Using non-JAX functions in JAX transformations
    This error can occur if you attempt to use a non-JAX library like ``numpy``
    or ``scipy`` inside a JAX transformation (:func:`~jax.jit`, :func:`~jax.grad`,
    :func:`jax.vmap`, etc.). For example::

      >>> from jax import jit
      >>> import numpy as np

      >>> @jit
      ... def func(x):
      ...   return np.sin(x)

      >>> func(np.arange(4))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerArrayConversionError: The numpy.ndarray conversion method
      __array__() was called on traced array with shape int32[4]

    In this case, you can fix the issue by using :func:`jax.numpy.sin` in place of
    :func:`numpy.sin`::

      >>> import jax.numpy as jnp
      >>> @jit
      ... def func(x):
      ...   return jnp.sin(x)

      >>> func(jnp.arange(4))
      Array([0.        , 0.84147096, 0.9092974 , 0.14112   ], dtype=float32)

    See also `External Callbacks`_ for options for calling back to host-side computations
    from transformed JAX code.

  Indexing a numpy array with a tracer
    If this error arises on a line that involves array indexing, it may be that
    the array being indexed ``x`` is a standard numpy.ndarray while the indices
    ``idx`` are traced JAX arrays. For example::

      >>> x = np.arange(10)

      >>> @jit
      ... def func(i):
      ...   return x[i]

      >>> func(0)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerArrayConversionError: The numpy.ndarray conversion method
      __array__() was called on traced array with shape int32[0]

    Depending on the context, you may fix this by converting the numpy array
    into a JAX array::

      >>> @jit
      ... def func(i):
      ...   return jnp.asarray(x)[i]

      >>> func(0)
      Array(0, dtype=int32)

    or by declaring the index as a static argument::

      >>> from functools import partial
      >>> @partial(jit, static_argnums=(0,))
      ... def func(i):
      ...   return x[i]

      >>> func(0)
      Array(0, dtype=int32)

  To understand more subtleties having to do with tracers vs. regular values,
  and concrete vs. abstract values, you may want to read
  :ref:`faq-different-kinds-of-jax-values`.

  .. _External Callbacks: https://docs.jax.dev/en/latest/notebooks/external_callbacks.html


  This error can occur when a JAX Tracer object is used in a context where a
  Python integer is expected (see :ref:`faq-different-kinds-of-jax-values` for
  more on what a Tracer is). It typically occurs in a few situations.

  Passing a tracer in place of an integer
    This error can occur if you attempt to pass a traced value to a function
    that requires a static integer argument; for example::

      >>> from jax import jit
      >>> import numpy as np

      >>> @jit
      ... def func(x, axis):
      ...   return np.split(x, 2, axis)

      >>> func(np.arange(4), 0)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerIntegerConversionError: The __index__() method was called on
      traced array with shape int32[0]

    When this happens, the solution is often to mark the problematic argument as
    static::

      >>> from functools import partial
      >>> @partial(jit, static_argnums=1)
      ... def func(x, axis):
      ...   return np.split(x, 2, axis)

      >>> func(np.arange(10), 0)
      [Array([0, 1, 2, 3, 4], dtype=int32),
       Array([5, 6, 7, 8, 9], dtype=int32)]

    An alternative is to apply the transformation to a closure that encapsulates
    the arguments to be protected, either manually as below or by using
    :func:`functools.partial`::

      >>> jit(lambda arr: np.split(arr, 2, 0))(np.arange(4))
      [Array([0, 1], dtype=int32), Array([2, 3], dtype=int32)]

    **Note a new closure is created at every invocation, which defeats the
    compilation caching mechanism, which is why static_argnums is preferred.**

  Indexing a list with a Tracer
    This error can occur if you attempt to index a Python list with a traced
    quantity.
    For example::

      >>> import jax.numpy as jnp
      >>> from jax import jit

      >>> L = [1, 2, 3]

      >>> @jit
      ... def func(i):
      ...   return L[i]

      >>> func(0)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerIntegerConversionError: The __index__() method was called on
      traced array with shape int32[0]

    Depending on the context, you can generally fix this either by converting
    the list to a JAX array::

      >>> @jit
      ... def func(i):
      ...   return jnp.array(L)[i]

      >>> func(0)
      Array(1, dtype=int32)

    or by declaring the index as a static argument::

      >>> from functools import partial
      >>> @partial(jit, static_argnums=0)
      ... def func(i):
      ...   return L[i]

      >>> func(0)
      Array(1, dtype=int32, weak_type=True)

  To understand more subtleties having to do with tracers vs. regular values,
  and concrete vs. abstract values, you may want to read
  :ref:`faq-different-kinds-of-jax-values`.


  This error occurs when a traced value in JAX is used in a context where a
  boolean value is expected (see :ref:`faq-different-kinds-of-jax-values`
  for more on what a Tracer is).

  The boolean cast may be an explicit (e.g. ``bool(x)``) or implicit, through use of
  control flow (e.g. ``if x > 0`` or ``while x``), use of Python boolean
  operators (e.g. ``z = x and y``, ``z = x or y``, ``z = not x``) or functions
  that use them (e.g. ``z = max(x, y)``, ``z = min(x, y)`` etc.).

  In some situations, this problem can be easily fixed by marking traced values as
  static; in others, it may indicate that your program is doing operations that are
  not directly supported by JAX's JIT compilation model.

  Examples:

  Traced value used in control flow
    One case where this often arises is when a traced value is used in
    Python control flow. For example::

      >>> from jax import jit
      >>> import jax.numpy as jnp
      >>> @jit
      ... def func(x, y):
      ...   return x if x.sum() < y.sum() else y

      >>> func(jnp.ones(4), jnp.zeros(4))  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerBoolConversionError: Attempted boolean conversion of JAX Tracer [...]

    We could mark both inputs ``x`` and ``y`` as static, but that would defeat
    the purpose of using :func:`jax.jit` here. Another option is to re-express
    the if statement in terms of the three-term :func:`jax.numpy.where`::

      >>> @jit
      ... def func(x, y):
      ...   return jnp.where(x.sum() < y.sum(), x, y)

      >>> func(jnp.ones(4), jnp.zeros(4))
      Array([0., 0., 0., 0.], dtype=float32)

    For more complicated control flow including loops, see
    :ref:`lax-control-flow`.

  Control flow on traced values
    Another common cause of this error is if you inadvertently trace over a boolean
    flag. For example::

      >>> @jit
      ... def func(x, normalize=True):
      ...   if normalize:
      ...     return x / x.sum()
      ...   return x

      >>> func(jnp.arange(5), True)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerBoolConversionError: Attempted boolean conversion of JAX Tracer ...

    Here because the flag ``normalize`` is traced, it cannot be used in Python
    control flow. In this situation, the best solution is probably to mark this
    value as static::

      >>> from functools import partial
      >>> @partial(jit, static_argnames=['normalize'])
      ... def func(x, normalize=True):
      ...   if normalize:
      ...     return x / x.sum()
      ...   return x

      >>> func(jnp.arange(5), True)
      Array([0. , 0.1, 0.2, 0.3, 0.4], dtype=float32)

    For more on ``static_argnums``, see the documentation of :func:`jax.jit`.

  Using non-JAX aware functions
    Another common cause of this error is using non-JAX aware functions within JAX
    code. For example:

      >>> @jit
      ... def func(x):
      ...   return min(x, 0)

      >>> func(2)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      TracerBoolConversionError: Attempted boolean conversion of JAX Tracer ...

    In this case, the error occurs because Python's built-in ``min`` function is not
    compatible with JAX transforms. This can be fixed by replacing it with
    ``jnp.minimum``:

      >>> @jit
      ... def func(x):
      ...   return jnp.minimum(x, 0)

      >>> print(func(2))
      0

  To understand more subtleties having to do with tracers vs. regular values,
  and concrete vs. abstract values, you may want to read
  :ref:`faq-different-kinds-of-jax-values`.


  This error occurs when you use a JAX value that has leaked out of a function.
  What does it mean to leak a value? If you use a JAX transformation on a
  function ``f`` that stores, in some scope outside of ``f``, a reference to
  an intermediate value, that value is considered to have been leaked.
  Leaking values is a side effect. (Read more about avoiding side effects in
  `Pure Functions <https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions>`_)

  JAX detects leaks when you then use the leaked value in another
  operation later on, at which point it raises an ``UnexpectedTracerError``.
  To fix this, avoid side effects: if a function computes a value needed
  in an outer scope, return that value from the transformed function explicitly.

  Specifically, a ``Tracer`` is JAX's internal representation of a function's
  intermediate values during transformations, e.g. within :func:`~jax.jit`,
  :func:`~jax.pmap`, :func:`~jax.vmap`, etc. Encountering a ``Tracer`` outside
  of a transformation implies a leak.

  Life-cycle of a leaked value
    Consider the following example of a transformed function which leaks a value
    to an outer scope::

      >>> from jax import jit
      >>> import jax.numpy as jnp

      >>> outs = []
      >>> @jit                   # 1
      ... def side_effecting(x):
      ...   y = x + 1            # 3
      ...   outs.append(y)       # 4

      >>> x = 1
      >>> side_effecting(x)      # 2
      >>> outs[0] + 1            # 5  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      UnexpectedTracerError: Encountered an unexpected tracer.

    In this example we leak a Traced value from an inner transformed scope to an
    outer scope. We get an ``UnexpectedTracerError`` when the leaked value is
    used, not when the value is leaked.

    This example also demonstrates the life-cycle of a leaked value:

      1. A function is transformed (in this case, by :func:`~jax.jit`)
      2. The transformed function is called (initiating an abstract trace of the
         function and turning ``x`` into a ``Tracer``)
      3. The intermediate value ``y``, which will later be leaked, is created
         (an intermediate value of a traced function is also a ``Tracer``)
      4. The value is leaked (appended to a list in an outer scope, escaping
         the function through a side-channel)
      5. The leaked value is used, and an UnexpectedTracerError is raised.

    The UnexpectedTracerError message tries to point to these locations in your
    code by including information about each stage. Respectively:

      1. The name of the transformed function (``side_effecting``) and which
         transform kicked off the trace  :func:`~jax.jit`).
      2. A reconstructed stack trace of where the leaked Tracer was created,
         which includes where the transformed function was called.
         (``When the Tracer was created, the final 5 stack frames were...``).
      3. From the reconstructed stack trace, the line of code that created
         the leaked Tracer.
      4. The leak location is not included in the error message because it is
         difficult to pin down! JAX can only tell you what the leaked value
         looks like (what shape it has and where it was created) and what
         boundary it was leaked over (the name of the transformation and the
         name of the transformed function).
      5. The current error's stack trace points to where the value is used.

    The error can be fixed by the returning the value out of the
    transformed function::

      >>> from jax import jit
      >>> import jax.numpy as jnp

      >>> outs = []
      >>> @jit
      ... def not_side_effecting(x):
      ...   y = x+1
      ...   return y

      >>> x = 1
      >>> y = not_side_effecting(x)
      >>> outs.append(y)
      >>> outs[0] + 1  # all good! no longer a leaked value.
      Array(3, dtype=int32, weak_type=True)

  Leak checker
    As discussed in point 2 and 3 above, JAX shows a reconstructed stack trace
    which points to where the leaked value was created.  This is because
    JAX only raises an error when the leaked value is used, not when the
    value is leaked. This is not the most useful place to raise this error,
    because you need to know the location where the Tracer was leaked to fix the
    error.

    To make this location easier to track down, you can use the leak checker.
    When the leak checker is enabled, an error is raised as soon as a ``Tracer``
    is leaked. (To be more exact, it will raise an error when the transformed
    function from which the ``Tracer`` is leaked returns)

    To enable the leak checker you can use the ``JAX_CHECK_TRACER_LEAKS``
    environment variable or the ``with jax.checking_leaks()`` context manager.

    .. note::
      Note that this tool is experimental and may report false positives. It
      works by disabling some JAX caches, so it will have a negative effect on
      performance and should only be used when debugging.

    Example usage::

      >>> from jax import jit
      >>> import jax.numpy as jnp

      >>> outs = []
      >>> @jit
      ... def side_effecting(x):
      ...   y = x+1
      ...   outs.append(y)

      >>> x = 1
      >>> with jax.checking_leaks():
      ...   y = side_effecting(x)  # doctest: +IGNORE_EXCEPTION_DETAIL
      Traceback (most recent call last):
          ...
      Exception: Leaked Trace


  This error occurs when a PRNG key is reused in an unsafe manner.
  Key reuse is checked only when `jax_debug_key_reuse` is
  set to `True`.

  Here is a simple example of code that would lead to such an error::

    >>> with jax.debug_key_reuse(True):  # doctest: +SKIP
    ...   key = jax.random.key(0)
    ...   value = jax.random.uniform(key)
    ...   new_value = jax.random.uniform(key)
    ...
    ---------------------------------------------------------------------------
    KeyReuseError                             Traceback (most recent call last)
    ...
    KeyReuseError: Previously-consumed key passed to jit-compiled function at index 0

  This sort of key reuse is problematic because the JAX PRNG is stateless, and keys
  must be manually split; For more information on this see `the Pseudorandom Numbers
  tutorial <https://docs.jax.dev/en/latest/random-numbers.html>`_.


See docstring for effects.py module for the calling convention for tokens.


Deferred call to `pxla.shard_args`.

  Per-array impls return this object instead of a result array to indicate a
  deferred `shard_args` call. `_batched_device_put_impl` then batches all
  `_DeferredShardArg` objects into a single `shard_args` call.


MappingsGenerator is a builder API for mappings.

  TC39 mapping data is inconvenient to emit directly: in an effort to compress
  data
  it encodes most indices using values _relative_ to the previous element.
  MappingsGenerator simplifies things by taking absolute indices everywhere.


Tuple describing how to partition an array across a mesh of devices.

  Each element is either ``None``, a string, or a tuple of strings.
  See the documentation of :class:`jax.sharding.NamedSharding` for more details.

  This class exists so JAX's pytree utilities can distinguish a partition
  specifications from tuples that should be treated as pytrees.


Inserts a CustomCallOp into the XLA graph with custom SPMD lowering rules.

  .. code-block:: python

    @custom_partitioning
    def f(*args):
      return ...

    def propagate_user_sharding(mesh, user_shape):
      '''Update the sharding of the op from a user's shape.sharding.'''
      user_sharding = jax.tree.map(lambda x: x.sharding, user_shape)

    def partition(mesh, arg_shapes, result_shape):
      def lower_fn(*args):
        ... builds computation on per-device shapes ...
      result_shardings = jax.tree.map(lambda x: x.sharding, result_shape)
      arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
      # result_sharding and arg_shardings may optionally be modified and the
      # partitioner will insert collectives to reshape.
      return mesh, lower_fn, result_sharding, arg_shardings

    def infer_sharding_from_operands(mesh, arg_shapes, shape):
      '''Compute the result sharding from the sharding of the operands.'''
      arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)

    f.def_partition(partition, propagate_user_sharding,
                    infer_sharding_from_operands=infer_sharding_from_operands,
                    sharding_rule='i j -> 'i j')

  The args to ``def_partition`` are as follows:

  * ``propagate_user_sharding``: Callable which takes the sharding of a user (in the dag)
    and returns a suggestion for a new `NamedSharding`. The default value is None.
    A trivial implementation is just to return the input sharding.
  * ``partition``: Callable which takes the SPMD suggested partition shapes and
    partition specs and returns the mesh, a per-shard lowering function, and the final
    input and output sharding specs (the SPMD partitioner will repartition the
    inputs to match). The mesh is returned to allow configuring axis_names for
    collectives when no mesh is provided.
  * ``infer_sharding_from_operands``: Callable which computes an output ``NamedSharding``
    from the ``NamedSharding`` chosen for each argument.
  * ``decode_shardings``: When set to True, convert input ``GSPMDSharding``s to
    ``NamedSharding`` if possible. This may not be possible if the user does not
    provide a contextual mesh.
  * ``sharding_rule``: an SdyShardingRule object, an Einsum-like notation string
    that describes the sharding rule, or a Callable that produces either of
    these. We borrow the idea from the einops.rearrange string , to use a space
    separator between factors and allow multiple letters factor names. See
    `jax-shardy-guide <https://colab.sandbox.google.com/github/openxla/shardy/blob/main/docs/getting_started_jax.ipynb>`_
    for more details and examples on how to use this.

  When config.use_shardy_partitioner.value is True, `sharding_rule` is used;
  otherwise, `propagate_user_sharding` and `infer_sharding_from_operands` are
  used.

  Positional arguments can be specified as static using static_argnums. JAX uses
  :code:`inspect.signature(fun)` to resolve these positional arguments.

  Examples:

    As an example, assume we want to enhance the existing ``jax.numpy.fft.fft``. This function computes
    the discrete Fourier transform of an N-dimensional input along the last dimension, and is batched
    along the first N-1 dimensions.
    By default, however, it will ignore the sharding of the input and gather the input on all devices.
    However, since ``jax.numpy.fft.fft`` is batched along the first N-1 dimensions,
    this is unnecessary. We will create a new ``my_fft`` op that, instead, does not alter the sharding
    along the first `N-1` dimensions, and only gathers the input along the last dimension if needed.

    .. code-block:: python

      import jax
      from jax.sharding import NamedSharding
      from jax.experimental.custom_partitioning import custom_partitioning
      from jax.experimental.pjit import pjit
      from jax.sharding import PartitionSpec as P
      from jax.sharding import Mesh
      from jax.numpy.fft import fft
      import regex as re
      import numpy as np

      # Pattern to detect all-gather or dynamic-slice in the generated HLO
      _PATTERN = '(dynamic-slice|all-gather)'

      # For an N-D input, keeps sharding along the first N-1 dimensions
      # but replicate along the last dimension
      def supported_sharding(sharding, shape):
          rank = len(shape.shape)
          max_shared_dims = min(len(sharding.spec), rank-1)
          names = tuple(sharding.spec[:max_shared_dims]) + tuple(None for _ in range(rank - max_shared_dims))
          return NamedSharding(sharding.mesh, P(*names))

      def partition(mesh, arg_shapes, result_shape):
          result_shardings = jax.tree.map(lambda x: x.sharding, result_shape)
          arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
          return mesh, fft, \
              supported_sharding(arg_shardings[0], arg_shapes[0]), \
              (supported_sharding(arg_shardings[0], arg_shapes[0]),)

      def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
          arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
          return supported_sharding(arg_shardings[0], arg_shapes[0])

      @custom_partitioning
      def my_fft(x):
          return fft(x)

      # Use Einsum-like notation to specify the sharding rule.
      my_fft.def_partition(
        infer_sharding_from_operands=infer_sharding_from_operands,
        partition=partition,
        sharding_rule='...i -> ...i')
      # Use SdyShardingRule object to specify the sharding rule.
      my_fft.def_partition(
        infer_sharding_from_operands=infer_sharding_from_operands,
        partition=partition,
        sharding_rule=SdyShardingRule(operand_mappings=((SDY_BATCHING, 'i'),), result_mappings=((SDY_BATCHING, 'i'),))))

    Now create a 2D array sharded along the first axis, pass it through ``my_fft``
    and notice how it is still sharded as expected, and identical to the output
    of ``fft``. However, inspecting the HLO
    (using ``lower(x).compile().runtime_executable().hlo_modules()``) reveals that
    ``my_fft`` does not create any all-gather or dynamic-slice, while ``fft`` does.

    .. code-block::

      with Mesh(np.array(jax.devices()), ('x',)):
        x = np.asarray(np.random.randn(32*1024, 1024), dtype=np.complex64)
        y = pjit(lambda x: x, in_shardings=None, out_shardings=P('x'))(x)
        pjit_my_fft = pjit(my_fft, in_shardings=P('x'), out_shardings=P('x'))
        pjit_fft    = pjit(fft,    in_shardings=P('x'), out_shardings=P('x'))
        print(pjit_my_fft(y))
        print(pjit_fft(y))
        # dynamic-slice or all-gather are not present in the HLO for my_fft, because x is a 2D array
        assert(re.search(_PATTERN, pjit_my_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string()) is None)
        # dynamic-slice or all-gather are present in the HLO for fft
        assert(re.search(_PATTERN, pjit_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string())    is not None)

    .. code-block::

      # my_fft
      [[-38.840824   +0.j        -40.649452  +11.845365j
      ...
        -1.6937828  +0.8402481j  15.999859   -4.0156755j]]

      # jax.numpy.fft.fft
      [[-38.840824   +0.j        -40.649452  +11.845365j
        ...
        -1.6937828  +0.8402481j  15.999859   -4.0156755j]]

    Because of the logic in ``supported_sharding``, ``my_fft`` also works on 1-dimensional arrays.
    However, in this case, the HLO of ``my_fft`` does show a dynamic-slice, since the last dimension
    is the dimension along which FFTs are calculated and needs to be replicated on all devices before
    the computation can be done.

    .. code-block::

      with Mesh(np.array(jax.devices()), ('x',)):
        x = np.asarray(np.random.randn(32*1024*1024), dtype=np.complex64)
        y = pjit(lambda x: x, in_shardings=None, out_shardings=P('x'))(x)
        pjit_my_fft = pjit(my_fft, in_shardings=P('x'), out_shardings=P('x'))
        pjit_fft    = pjit(fft,    in_shardings=P('x'), out_shardings=P('x'))
        print(pjit_my_fft(y))
        print(pjit_fft(y))
        # dynamic-slice or all-gather are present in the HLO for my_fft, because x is a 1D array
        assert(re.search(_PATTERN, pjit_my_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string()) is None)
        # dynamic-slice or all-gather are present in the HLO for fft
        assert(re.search(_PATTERN, pjit_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string())    is not None)

    .. code-block::

      # my_fft
      [    7.217285   +0.j     -3012.4937  +4287.635j   -405.83594 +3042.984j
      ...  1422.4502  +7271.4297j  -405.84033 -3042.983j
      -3012.4963  -4287.6343j]

      # jax.numpy.fft.fft
      [    7.217285   +0.j     -3012.4937  +4287.635j   -405.83594 +3042.984j
      ...  1422.4502  +7271.4297j  -405.84033 -3042.983j
      -3012.4963  -4287.6343j]


Add a signal handler that sets a flag if SIGINT or SIGTERM are caught.


Set up a JAX-transformable function for a custom JVP rule definition.

  This class is meant to be used as a function decorator. Instances are
  callables that behave similarly to the underlying function to which the
  decorator was applied, except when a differentiation transformation (like
  :py:func:`jax.jvp` or :py:func:`jax.grad`) is applied, in which case a custom
  user-supplied JVP rule function is used instead of tracing into and
  performing automatic differentiation of the underlying function's
  implementation.

  There are two instance methods available for defining the custom JVP rule:
  :py:func:`~jax.custom_jvp.defjvp` for defining a *single* custom JVP rule for
  all the function's inputs, and for convenience
  :py:func:`~jax.custom_jvp.defjvps`, which wraps
  :py:func:`~jax.custom_jvp.defjvp`, and allows you to provide separate
  definitions for the partial derivatives of the function w.r.t. each of its
  arguments.

  For example::

    @jax.custom_jvp
    def f(x, y):
      return jnp.sin(x) * y

    @f.defjvp
    def f_jvp(primals, tangents):
      x, y = primals
      x_dot, y_dot = tangents
      primal_out = f(x, y)
      tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
      return primal_out, tangent_out

  For a more detailed introduction, see the tutorial_.

  .. _tutorial: https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html


Set up a JAX-transformable function for a custom VJP rule definition.

  This class is meant to be used as a function decorator. Instances are
  callables that behave similarly to the underlying function to which the
  decorator was applied, except when a reverse-mode differentiation
  transformation (like :py:func:`jax.grad`) is applied, in which case a custom
  user-supplied VJP rule function is used instead of tracing into and performing
  automatic differentiation of the underlying function's implementation. There
  is a single instance method, :py:func:`~jax.custom_vjp.defvjp`, which may be
  used to define the custom VJP rule.

  This decorator precludes the use of forward-mode automatic differentiation.

  For example::

    @jax.custom_vjp
    def f(x, y):
      return jnp.sin(x) * y

    def f_fwd(x, y):
      return f(x, y), (jnp.cos(x), jnp.sin(x), y)

    def f_bwd(res, g):
      cos_x, sin_x, y = res
      return (cos_x * g * y, sin_x * g)

    f.defvjp(f_fwd, f_bwd)

  For a more detailed introduction, see the tutorial_.

  .. _tutorial: https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html


Specifies a PRNG key implementation.


Handles requests from `ui.perfetto.dev` for the `trace.json`


Context manager that generates a trace event in the profiler.

  The trace event spans the duration of the code enclosed by the context.

  For example:

  >>> x = jnp.ones((1000, 1000))
  >>> with jax.profiler.TraceAnnotation("my_label"):
  ...   result = jnp.dot(x, x.T).block_until_ready()

  This will cause a "my_label" event to show up on the trace timeline if the
  event occurs while the process is being traced.


Context manager that generates a step trace event in the profiler.

  The step trace event spans the duration of the code enclosed by the context.
  The profiler will provide the performance analysis for each step trace event.

  For example, it can be used to mark training steps and enable the profiler to
  provide the performance analysis per step:

  >>> while global_step < NUM_STEPS:                                           # doctest: +SKIP
  ...   with jax.profiler.StepTraceAnnotation("train", step_num=global_step):  # doctest: +SKIP
  ...     train_step()                                                         # doctest: +SKIP
  ...     global_step += 1                                                     # doctest: +SKIP

  This will cause a "train xx" event to show up on the trace timeline if the
  event occurs while the process is being traced by TensorBoard. In addition,
  if using accelerators, the device trace timeline will also show a "train xx"
  event. Note that "step_num" can be set as a keyword argument to pass the
  global step number to the profiler.


Declare the hardware resources available in the scope of this manager.

  See `Distributed arrays and automatic parallelization`_ and
  `Explicit Sharding`_ tutorials.

  Args:
    devices: A NumPy ndarray object containing JAX device objects (as
      obtained e.g. from :py:func:`jax.devices`).
    axis_names: A sequence of resource axis names to be assigned to the
      dimensions of the ``devices`` argument. Its length should match the
      rank of ``devices``.
    axis_types: and optional tuple of :class:`jax.sharding.AxisType` entries corresponding to
      the ``axis_names``. See `Explicit Sharding`_ for more information.

  Examples:

    >>> from jax.sharding import Mesh
    >>> from jax.sharding import PartitionSpec as P, NamedSharding
    >>> import numpy as np
    ...
    >>> # Declare a 2D mesh with axes `x` and `y`.
    >>> devices = np.array(jax.devices()).reshape(4, 2)
    >>> mesh = Mesh(devices, ('x', 'y'))
    >>> inp = np.arange(16).reshape(8, 2)
    >>> arr = jax.device_put(inp, NamedSharding(mesh, P('x', 'y')))
    >>> out = jax.jit(lambda x: x * 2)(arr)
    >>> assert out.sharding == NamedSharding(mesh, P('x', 'y'))

  .. _Distributed arrays and automatic parallelization: https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
  .. _Explicit Sharding:  https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html


AbstractMesh contains only axis names and axis sizes.

  It does not contain concrete devices compared to `jax.sharding.Mesh`. You
  should use this as an input to the sharding passed to with_sharding_constraint
  and mesh passed to shard_map to avoid tracing and lowering cache misses when
  your mesh shape and axis names stay the same but the devices change.
  See the description of https://github.com/jax-ml/jax/pull/23022 for more
  details.

  Args:
    axis_sizes: A tuple of integers specifying the size of each resource axis.
    axis_names: A tuple of resource axis names to be assigned to the
      dimensions of the ``devices`` argument. Its length should match the
      rank of ``devices``.
    axis_types: and optional tuple of :class:`jax.sharding.AxisType` entries corresponding to
      the ``axis_names``. See `Explicit Sharding`_ for more information.

  .. _Explicit Sharding:  https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html


Array base class for JAX

  ``jax.Array`` is the public interface for instance checks and type annotation
  of JAX arrays and tracers. Its main applications are in instance checks and
  type annotations; for example::

    x = jnp.arange(5)
    isinstance(x, jax.Array)  # returns True both inside and outside traced functions.

    def f(x: Array) -> Array:  # type annotations are valid for traced and non-traced types.
      return x

  ``jax.Array`` should not be used directly for creation of arrays; instead you
  should use array creation routines offered in :mod:`jax.numpy`, such as
  :func:`jax.numpy.array`, :func:`jax.numpy.zeros`, :func:`jax.numpy.ones`,
  :func:`jax.numpy.full`, :func:`jax.numpy.arange`, etc.


Decouples function equality and hash from its identity.

  Local lambdas and function defs are reallocated on each function call, making
  the functions created on different calls compare as unequal. This breaks our
  caching logic, which should really only care about comparing the semantics and
  not actual identity.

  This class makes it possible to compare different functions based on their
  semantics. The parts that are taken into account are: the bytecode of the
  wrapped function (which is cached by the CPython interpreter and is stable
  across the invocations of the surrounding function), and `closure` which
  should contain all values in scope that affect the function semantics. In
  particular `closure` should contain all elements of the function closure, or
  it should be possible to derive the relevant elements of the true function
  closure based solely on the contents of the `closure` argument (e.g. in case
  some closed-over values are not hashable, but are entirely determined by
  hashable locals).


A variant of `abc.ABCMeta` which does not allow virtual subclasses.

  Virtual subclasses support require `abc.ABCMeta` to roundtrip through
  pure Python when doing instance/subclass checking. This if fine for ABCs
  which need virtual subclasses, but is wasteful for the ones which don't.


Python exception which can contain an error message with JAX run-time info.


Customize the DCE behavior of a JAX-transformable function.

  JAX uses dead code elimination (DCE) to remove unused computations from a
  JAX program. This typically works transparently when the program is
  completely specified by known JAX operations, but opaque kernels like calls
  to :py:func:`~jax.experimental.pallas.pallas_call` or
  :py:func:`~jax.ffi.ffi_call`, for example, may cause problems.

  In JAX, DCE is performed when a function is staged out using
  :py:func:`jax.jit`, so it won't be applied when running JAX in eager mode.
  Similarly, the ``custom_dce`` decorator requires that both the decorated
  function and the custom DCE rule be compatible with :py:func:`~jax.jit`.

  This decorator allows users to customize the DCE behavior of a function by
  defining a custom DCE rule. For a ``custom_dce`` wrapped function
  ``f(*args)``, the signature of the DCE rule is ``dce_rule(used_outs, *args)``
  where ``used_outs`` is a Pytree with the same structure as the output of
  ``f``, and each leaf is is a ``bool`` indicating which outputs should
  be computed. The remaining arguments ``*args`` are the original arguments to
  ``f``. The rule ``dce_rule`` should return a Pytree with the same structure
  as the original output of ``f``, but any unused outputs can be replaced with
  ``None``.

  For example::

    >>> @jax.experimental.custom_dce.custom_dce
    ... def f(x, y):
    ...   return jnp.sin(x) * y, x * jnp.sin(y)
    ...
    >>> @f.def_dce
    ... def f_dce_rule(used_outs, x, y):
    ...   return (
    ...       jnp.sin(x) * y if used_outs[0] else None,
    ...       x * jnp.sin(y) if used_outs[1] else None,
    ...   )

  In this example, ``used_outs`` is a ``tuple`` with two ``bool`` values,
  indicating which outputs are required. The DCE rule only computes the
  required outputs, replacing the unused outputs with ``None``.

  If the ``static_argnums`` argument is provided to ``custom_dce``, the
  indicated arguments are treated as static when the function is traced, and
  they will be moved to the front when calling the DCE rule. For example, if
  ``fun`` takes 2 arguments ``fun(x, y)``, and ``static_argnums`` is ``(1,)``,
  then the DCE rule will be called as ``dce_rule(y, used_outs, x)``.


Describes how a :class:`jax.Array` is laid out across devices.


Effort level enum, mirroring the XLA effort options.


A holder for a configuration value.

  There are two kinds of value holders: ``Flag``, which is assigned exactly
  once and never modified after; and ``State``, which can be changed locally
  within a thread via a context manager.


Scalar class for extended dtypes.

  This is an abstract class that should never be instantiated, but rather
  exists for the sake of `jnp.issubdtype`.

  Examples:
    >>> from jax import random
    >>> from jax import dtypes
    >>> key = random.key(0)
    >>> jnp.issubdtype(key.dtype, dtypes.extended)
    True


Scalar class for PRNG Key dtypes.

  This is an abstract class that should never be instantiated, but rather
  exists for the sake of `jnp.issubdtype`.

  Examples:
    >>> from jax import random
    >>> from jax import dtypes
    >>> key = random.key(0)
    >>> jnp.issubdtype(key.dtype, dtypes.prng_key)
    True


Abstract Base Class for extended dtypes


Object that delegates __call__, __hash__, and __eq__ to another object.


A version of functools.partial that works in pytrees.

  Use it for partial function evaluation in a way that is compatible with JAX's
  transformations, e.g., ``Partial(func, *args, **kwargs)``.

  (You need to explicitly opt-in to this behavior because we didn't want to give
  functools.partial different semantics than normal function closures.)

  For example, here is a basic usage of ``Partial`` in a manner similar to
  ``functools.partial``:

  >>> import jax.numpy as jnp
  >>> add_one = Partial(jnp.add, 1)
  >>> add_one(2)
  Array(3, dtype=int32, weak_type=True)

  Pytree compatibility means that the resulting partial function can be passed
  as an argument within transformed JAX functions, which is not possible with a
  standard ``functools.partial`` function:

  >>> from jax import jit
  >>> @jit
  ... def call_func(f, *args):
  ...   return f(*args)
  ...
  >>> call_func(add_one, 2)
  Array(3, dtype=int32, weak_type=True)

  Passing zero arguments to ``Partial`` effectively wraps the original function,
  making it a valid argument in JAX transformed functions:

  >>> call_func(Partial(jnp.add), 1, 2)
  Array(3, dtype=int32, weak_type=True)

  Had we passed ``jnp.add`` to ``call_func`` directly, it would have resulted in
  a ``TypeError``.

  Note that if the result of ``Partial`` is used in the context where the
  value is traced, it results in all bound arguments being traced when passed
  to the partially-evaluated function:

  >>> print_zero = Partial(print, 0)
  >>> print_zero()
  0
  >>> call_func(print_zero)  # doctest:+ELLIPSIS
  JitTracer<~int32[]>


  Wraps a TestResult to make it thread safe.

  We do this by accumulating API calls and applying them in a batch under a
  lock at the conclusion of each test case.

  We duck type instead of inheriting from TestResult because we aren't actually
  a perfect implementation of TestResult, and would rather get a loud error
  for things we haven't implemented.


Runs tests in parallel using threads if TEST_NUM_THREADS is > 1.

  Caution: this test suite does not run setUpClass or setUpModule methods if
  thread parallelism is enabled.


Thread-local state that contains a list of warning handlers.


Describes the factors for a compound factor.

  A compound factor should contain at least two factors, e.g.
  * CompoundFactor('b', 'c').


Describes the factors for an operand or result.

  Each element is either a factor or a CompoundFactor. A leading element can
  also be BATCHING, which represents batching dimensions. examples:
  * ArrayMapping('a')
  * ArrayMapping('b', 'c')
  * ArrayMapping(CompoundFactor('b', 'c'), 'd')
  * ArrayMapping(BATCHING, CompoundFactor('b', 'c'), 'd')


Represents a Shardy sharding rule.

  An SdyShardingRule contains the ArrayMappings for operands and results, and an
  optional list of factor sizes. A factor is a name used in the ArrayMappings.
  If a factor is only used in CompoundFactors, its size must be specified.


Bounded cache with least-recently-used (LRU) eviction policy.

  This implementation includes cache reading, writing and eviction
  based on the LRU policy.

  Notably, when ``max_size`` is set to -1, the cache eviction
  is disabled, and the LRU cache functions as a normal cache
  without any size limitations.


A side-effect associated with the input of a `JaxprEqn` or a `Jaxpr`.

  This is used as a base class for effects associated with inputs, e.g.,
  reading/writing from mutable inputs.

  When used in a `JaxprEqn`, `input_index` refers to `eqn.invars`.
  When used in a `Jaxpr`, `input_index` refers to `jaxpr.constvars + jaxpr.invars`.


Box object used when comparing static arguments as a jit key.

  Requires exact type equality using `is` and value equality.


A Python function callable with flat arguments and results.

  An instance of this class is used as a parameter for the callback primitives.
  We prefer it to an anonymous flattened function because it produces
  equal objects when we call the same Python function with the same argument
  structure.


Compiled representation of a function specialized to types/values.

  A compiled computation is associated with an executable and the
  remaining information needed to execute it. It also provides a
  common API for querying properties of compiled computations across
  JAX's various compilation paths and backends.


Lowering of a function specialized to argument types and values.

  A lowering is a computation ready for compilation. This class
  carries a lowering together with the remaining information needed to
  later compile and execute it. It also provides a common API for
  querying properties of lowered computations across JAX's various
  lowering paths (:func:`~jax.jit`, :func:`~jax.pmap`, etc.).


Traced form of a function specialized to argument types and values.

  A traced computation is ready for lowering. This class carries the
  traced representation with the remaining information needed to later
  lower, compile, and execute it.


A function ready to be traced, lowered, and compiled.

  This protocol reflects the output of functions such as
  ``jax.jit``. Calling it results in JIT (just-in-time) lowering,
  compilation, and execution. It can also be explicitly lowered prior
  to compilation, and the result compiled prior to execution.


Mutable array reference.

  In most cases this should not be constructed directly, but rather
  via :func:`jax.ref.array_ref`. For examples of how this can be
  used, refer to the `ArrayRef guide`_.

  .. _ArrayRef guide: https://docs.jax.dev/en/latest/array_refs.html


A container for the shape, dtype, and other static attributes of an array.

  ``ShapeDtypeStruct`` is often used in conjunction with :func:`jax.eval_shape`.

  Args:
    shape: a sequence of integers representing an array shape
    dtype: a dtype-like object
    sharding: (optional) a :class:`jax.Sharding` object


A :class:`Sharding` that places its data on a single device.

  Args:
    device: A single :py:class:`Device`.

  Examples:

    >>> single_device_sharding = jax.sharding.SingleDeviceSharding(
    ...     jax.devices()[0])


Describes a sharding used by :func:`jax.pmap`.


Represents a pmap mesh (only along the replica axes).


A hardware axis context for parallel computations that use the GSPMD partitioner.

  This includes the mesh that will later by used to execute this computation,
  as well as a set of mesh axes that are currently lowered in the MANUAL
  sharding mode.


A hardware axis context for parallel computations that are partitioned by JAX.

  Unlike in the SPMDAxisContext, this means that JAX might need to emit calls to
  explicit collectives.


A hardware axis context for parallel computations that use the sharding
  interface.

  This context also uses the GSPMD partitioner.


A decision procedure based on elimination of terms.

  Given an expression `e = t*t_k + rest_e` for which we want to compute bounds,
  and a constraint `c = t*t_c_k + rest_c >= 0`,

  Let `e0 = e*abs(t_c_k) - c*sgn(t_c_k)*t_k`. (Note that we eliminated `t` from
  `e0`, since `abs(t_c_k)*t_k = sgn(t_c_k)*t_k*t_c_k`.)

  Since `c >= 0`,
  if `sgn(t_c_k)*t_k > 0`:
     then `abs(t_c_k)*e >= e0`, hence, `LB(e) >= ceil(LB(e0) / abs(t_c_k))`,

  if `sgn(t_c_k)*t_k < 0`
    then `abs(t_c_k)*e <= e0`, hence, `UB(e) <= floor(UB(e0) / abs(t_c_k))`,

  See the implementation in self.combine_term_with_existing.
  Do not use the constructor directly, use the `build` static method.


Raised when we cannot conclusively compute with symbolic dimensions.


Represents a factor in a symbolic dimension expression.

  Factors are either variables, or expressions of the form floordiv(E1, E2) or
  mod(E1, E2), or max(E1, E2), or min(E1, E2).
  Factors are multiplied to form terms (see _DimTerm), and
  terms are added to form symbolic expressions (see _DimExpr).

  Args:
    * var: if specified then the factor is a dimension variable. `operation`
      must be `None`.
    * operation: if specified then the factor is an operation applied to
      `operands`. One of `FLOORDIR` or `MOD` or `MAX` or `MIN`.
      `var` must be `None`
    * operands: the operands to which the operation is applied.


Represents a multiplication of factors.

  The representation is a sequence of _DimFactor factors along with their
  integer exponents (>= 1). The empty sequence represents the constant 1.


Symbolic expressions using dimension variables.

  A dimension expression is an addition of terms (_DimTerm), which themselves
  are products of factors (_DimFactor).

  The representation of a _DimExpr is as sequence of pairs `(term, coeff)`,
  representing the linear combination of terms with the given coefficients.
  The sequence is sorted by lexicographic (syntactic) ordering of `_DimTerm`,
  with the largest terms first. The special term `_DimTerm_one` is mapped
  to the free integer coefficient of the expression.

  We overload integer operations, but we do that soundly, raising
  :class:`InconclusiveDimensionOperation` when the result is not
  representable as a _DimExpr.


Identifies a scope for symbolic expressions.

  All symbolic expressions that interact (e.g., appear in the argument shapes
  for one JAX function invocation, or are involved in arithmetic operations)
  must be from the same scope and must share the same SymbolicScope object.

  Holds the constraints on symbolic expressions.

  See [the README](https://docs.jax.dev/en/latest/export/shape_poly.html#user-specified-symbolic-constraints)
  for more details.

  Args:
    constraints_str: A sequence of constraints on symbolic dimension expressions,
      of the form `e1 >= e2` or `e1 <= e2` or `e1 == e2`.


Specifies desired precision for the bounds calculation.

  Since the bounds calculations are expensive, we allow the caller to specify
  a sufficient condition for a result. As the bounds calculations progresses
  the lower bounds is progressively increased and the upper bounds is
  progressively decreased. Depending on the precision, we may stop the
  computation early, if the results are sufficient for the use case.

  The enumeration values are chosen such that, if "(lb, ub)" are sufficient
  for a precision value then they are also sufficient for any smaller
  precision.


Tuple of polymorphic dimension specifications.

  See docstring of :func:`jax2tf.convert`.


A safety check that should be skipped on (de)serialization.

  Most of these checks are performed on serialization, but some are deferred to
  deserialization. The list of disabled checks is attached to the serialization,
  e.g., as a sequence of string attributes to `jax.export.Exported` or of
  `tf.XlaCallModuleOp`.

  When using jax2tf, you can disable more deserialization safety checks
  by passing `TF_XLA_FLAGS=--tf_xla_call_module_disabled_checks=platform`.


A JAX function lowered to StableHLO.

  Attributes:
    fun_name: the name of the exported function, for error messages.
    in_tree: a PyTreeDef describing the tuple (args, kwargs) of the lowered JAX
        function. The actual lowering does not depend on the `in_tree`, but this
        can be used to invoke the exported function using the same argument
        structure.
    in_avals: the flat tuple of input abstract values. May contain dimension
        expressions in the shapes.
    out_tree: a PyTreeDef describing the result of the lowered JAX function.
    out_avals: the flat tuple of output abstract values. May contain dimension
        expressions in the shapes, with dimension variables among those in
        `in_avals`.
    in_shardings_hlo: the flattened input shardings, a sequence as long
        as `in_avals`. `None` means unspecified sharding.
        Note that these do not include the mesh or the actual devices used in
        the mesh. See `in_shardings_jax` for a way to turn these
        into sharding specification that can be used with JAX APIs.
    out_shardings_hlo: the flattened output shardings, a sequence as long
        as `out_avals`. `None` means unspecified sharding.
        Note that these do not include the mesh or the actual devices used in
        the mesh. See `out_shardings_jax` for a way to turn these
        into sharding specification that can be used with JAX APIs.
    nr_devices: the number of devices that the module has been lowered for.
    platforms: a tuple containing the platforms for which the function should
        be exported. The set of platforms in JAX is open-ended; users can
        add platforms. JAX built-in platforms are: 'tpu', 'cpu', 'cuda', 'rocm'.
        See https://docs.jax.dev/en/latest/export/export.html#cross-platform-and-multi-platform-export.
    ordered_effects: the ordered effects present in the serialized module.
        This is present from serialization version 9. See https://docs.jax.dev/en/latest/export/export.html#module-calling-convention
        for the calling convention in presence of ordered effects.
    unordered_effects: the unordered effects present in the serialized module.
        This is present from serialization version 9.
    mlir_module_serialized: the serialized lowered VHLO module.
    calling_convention_version: a version number for the calling
        convention of the exported module.
        See more versioning details at https://docs.jax.dev/en/latest/export/export.html#calling-convention-versions.
    module_kept_var_idx: the sorted indices of the arguments among `in_avals` that
        must be passed to the module. The other arguments have been dropped
        because they are not used.
    uses_global_constants: whether the `mlir_module_serialized` uses shape
        polymorphism or multi-platform export.
        This may be because `in_avals` contains dimension
        variables, or due to inner calls of Exported modules that have
        dimension variables or platform index arguments. Such modules need
        shape refinement before XLA compilation.
    disabled_safety_checks: a list of descriptors of safety checks that have been
        disabled at export time. See docstring for `DisabledSafetyCheck`.
    _get_vjp: an optional function that takes the current exported function and
        returns the exported VJP function.
        The VJP function takes a flat list of arguments,
        starting with the primal arguments and followed by a cotangent argument
        for each primal output. It returns a tuple with the cotangents
        corresponding to the flattened primal inputs.

  See a [description of the calling convention for the `mlir_module`](https://docs.jax.dev/en/latest/export/export.html#module-calling-convention).


Abstract cluster supports both single and multislice TPU environments.

  If MEGASCALE_COORDINATOR_ADDRESS is not set, we assume single slice topology.
  Concrete extensions of this class must implement methods for generating a list
    of within-slice workers and a within-slice process ID.
  `get_coordinator_address` must return the address of the host with
  process ID 0 (as returned by `get_process_id`), since the coordinator service
  is started on the host with process ID = 0.


Interface for defining a cluster environment.

  To enable auto bootstrapping (aka :func:`jax.distributed.initialize()`),
  cluster environments need to derive from :class:`ClusterEnv` and implement
  :func:`is_env_present`, :func:`get_coordinator_address`,
  :func:`get_process_count`, and :func:`get_process_id`.
  :class:`ClusterEnv` subclasses are automatically detected when imported.


Rotation in 3 dimensions.

  JAX implementation of :class:`scipy.spatial.transform.Rotation`.

  Examples:
    Construct an object describing a 90 degree rotation about the z-axis:

    >>> from jax.scipy.spatial.transform import Rotation
    >>> r = Rotation.from_euler('z', 90, degrees=True)

    Convert to a rotation vector:

    >>> r.as_rotvec()
    Array([0.       , 0.       , 1.5707964], dtype=float32)

    Convert to rotation matrix:

    >>> r.as_matrix()
    Array([[ 0.        , -0.99999994,  0.        ],
           [ 0.99999994,  0.        ,  0.        ],
           [ 0.        ,  0.        ,  0.99999994]], dtype=float32)

    Compose with another rotation:

    >>> r2 = Rotation.from_euler('x', 90, degrees=True)
    >>> r3 = r * r2
    >>> r3.as_matrix()
    Array([[0., 0., 1.],
           [1., 0., 0.],
           [0., 1., 0.]], dtype=float32)

    See the scipy :class:`~scipy.spatial.transform.Rotation` documentation for
    further examples of manipulating Rotation objects.


Spherical Linear Interpolation of Rotations.

  JAX implementation of :class:`scipy.spatial.transform.Slerp`.

  Examples:
    Create a Slerp instance from a series of rotations:

    >>> import math
    >>> from jax.scipy.spatial.transform import Rotation, Slerp
    >>> rots = jnp.array([[90, 0, 0],
    ...                   [0, 45, 0],
    ...                   [0, 0, -30]])
    >>> key_rotations = Rotation.from_euler('zxy', rots, degrees=True)
    >>> key_times = [0, 1, 2]
    >>> slerp = Slerp.init(key_times, key_rotations)
    >>> times = [0, 0.5, 1, 1.5, 2]
    >>> interp_rots = slerp(times)
    >>> interp_rots.as_euler('zxy')
    Array([[ 1.5707963e+00,  0.0000000e+00,  0.0000000e+00],
           [ 8.5309029e-01,  3.8711953e-01,  1.7768645e-01],
           [-2.3841858e-07,  7.8539824e-01,  0.0000000e+00],
           [-5.6668043e-02,  3.9213133e-01, -2.8347540e-01],
           [ 0.0000000e+00,  0.0000000e+00, -5.2359891e-01]], dtype=float32)


Results from L-BFGS optimization

  Parameters:
    converged: True if minimization converged
    failed: True if non-zero status and not converged
    k: integer number of iterations of the main loop (optimisation steps)
    nfev: integer total number of objective evaluations performed.
    ngev: integer total number of jacobian evaluations
    x_k: array containing the last argument value found during the search. If
      the search converged, then this value is the argmin of the objective
      function.
    f_k: array containing the value of the objective function at `x_k`. If the
      search converged, then this is the (local) minimum of the objective
      function.
    g_k: array containing the gradient of the objective function at `x_k`. If
      the search converged the l2-norm of this tensor should be below the
      tolerance.
    status: integer describing the status:
      0 = nominal  ,  1 = max iters reached  ,  2 = max fun evals reached
      3 = max grad evals reached  ,  4 = insufficient progress (ftol)
      5 = line search failed
    ls_status: integer describing the end status of the last line search


Results from BFGS optimization.

  Parameters:
    converged: True if minimization converged.
    failed: True if line search failed.
    k: integer the number of iterations of the BFGS update.
    nfev: integer total number of objective evaluations performed.
    ngev: integer total number of jacobian evaluations
    nhev: integer total number of hessian evaluations
    x_k: array containing the last argument value found during the search. If
      the search converged, then this value is the argmin of the objective
      function.
    f_k: array containing the value of the objective function at `x_k`. If the
      search converged, then this is the (local) minimum of the objective
      function.
    g_k: array containing the gradient of the objective function at `x_k`. If
      the search converged the l2-norm of this tensor should be below the
      tolerance.
    H_k: array containing the inverse of the estimated Hessian.
    status: int describing end state.
    line_search_status: int describing line search end state (only means
      something if line search fails).


Object holding optimization results.

  Parameters:
    x: final solution.
    success: ``True`` if optimization succeeded.
    status: integer solver specific return code. 0 means converged (nominal),
      1=max BFGS iters reached, 3=zoom failed, 4=saddle point reached,
      5=max line search iters reached, -1=undefined
    fun: final function value.
    jac: final jacobian array.
    hess_inv: final inverse Hessian estimate.
    nfev: integer number of function calls used.
    njev: integer number of gradient evaluations.
    nit: integer number of iterations of the optimization algorithm.


Results of line search.

  Parameters:
    failed: True if the strong Wolfe criteria were satisfied
    nit: integer number of iterations
    nfev: integer number of functions evaluations
    ngev: integer number of gradients evaluations
    k: integer number of iterations
    a_k: integer step size
    f_k: final function value
    g_k: final gradient value
    status: integer end status


Gaussian Kernel Density Estimator

  JAX implementation of :class:`scipy.stats.gaussian_kde`.

  Parameters:
    dataset: arraylike, real-valued. Data from which to estimate the distribution.
      If 1D, shape is (n_data,). If 2D, shape is (n_dimensions, n_data).
    bw_method: string, scalar, or callable. Either "scott", "silverman", a scalar
      value, or a callable function which takes ``self`` as a parameter.
    weights: arraylike, optional. Weights of the same shape as the dataset.


A web-based debugger.


A mutable DOM element that displays code as HTML.


Displays information about a stack frame.


Main view for the Colab debugger.


A JAX debugger for a Colab environment.


A text-based debugger.


A DOM element that can be mutated.


A `div` that can be edited.


An immutable DOM element.


Encapsulates Python frame information.


Abstract mutable array reference.

  Refer to the `ArrayRef guide`_ for more information.

  .. _ArrayRef guide: https://docs.jax.dev/en/latest/array_refs.html


A slice with a start index and a size.

  Both start index and size can either be static, i.e. known at tracing
  and compilation time, or dynamic.


A partial discharge rule.

  Exactly like a discharge rule only it accepts a `should_discharge`
  argument that indicates which inputs should be discharged and the
  return value returns a tuple of which the first element is the new
  inputs or none but only the ones that correspond to `True` entries
  in `should_charge`.


Module-wide context information for MLIR lowering.


Per-rule context information for MLIR lowering.


An immutable container of tokens to be used to lower effectful jaxprs. When lowering
  effectful jaxprs, we need to thread HLO tokens to sequence them. Each effect
  will need its own token that will be threaded in and out of the effectful
  primitives. A `TokenSet` encapsulates a set of HLO tokens that will be
  used by the lowering rules.


Hashable wrapper of core.Literal, used for deduplicating IR constants.


The logic to shard inputs, execute a replicated model, returning outputs.


Partial value: either a known value or an unknown (abstract) value.

  Represented as a pair `(aval_opt, const)` of one of two kinds:
  * `(None, <Constant>)` indicates a known value, where the constant is either a
    Tracer or satisfies `core.valid_jaxtype(const)`;
  * `(<AbstractValue>, None)` indicates an unknown value characterized by an
    abstract value.


Image resize method.

  Possible values are:

  NEAREST:
    Nearest-neighbor interpolation.

  LINEAR:
    `Linear interpolation`_.

  LANCZOS3:
    `Lanczos resampling`_, using a kernel of radius 3.

  LANCZOS5:
    `Lanczos resampling`_, using a kernel of radius 5.

  CUBIC:
    `Cubic interpolation`_, using the Keys cubic kernel.

  .. _Linear interpolation: https://en.wikipedia.org/wiki/Bilinear_interpolation
  .. _Cubic interpolation: https://en.wikipedia.org/wiki/Bicubic_interpolation
  .. _Lanczos resampling: https://en.wikipedia.org/wiki/Lanczos_resampling


Descriptor for a randomly generated argument.

  See description of `Harness`.


Descriptor for a static argument.

  See description of `Harness`.


Descriptor for a dynamic argument generated from a PRNG factory.

  See description of `Harness`.


Specifies inputs and callable for a test harness.

  See the module docstring for an introduction to harnesses.

  A harness is conceptually a callable and a list of arguments, that together
  exercise a use case. The harness can optionally have additional parameters
  that can be used by the test.

  The arguments are specified through argument descriptors. An argument
  descriptor can be:
    * a numeric value or ndarray, or
    * an instance of ``RandArg(shape, dtype)`` to be used with a PRNG to
    generate random tensor of the given shape and type, or
    * an instance of ``CustomArg(fun)`` to be used with a PRNG, or
    * an instance of ``StaticArg(value)``. Often these are the non-array
      arguments, e.g., a shape.

  The given callable will be passed one argument corresponding to each
  argument descriptor, e.g., `harness.fun(* harness.args_maker(rng))`.
  However, in many applications we only care about the non-static arguments.
  For that purpose, you can use `harness.dyn_fun(*
  harness.dyn_args_maked(rng))`,
  where `harness.dyn_fun` is `harness.fun` specialized to the static arguments.

  For example, a harness for ``lax.take(arr, indices, axis=None)`` may want
  to expose as external (non-static) argument the array and the indices, and
  keep the axis as a static argument (technically specializing the `take` to
  a axis):

    Harness(lax.slice_p,
            f"take_axis={axis}",
            lax.take,
            [RandArg((2, 4), np.float32), np.array([-1, 0, 1]),
            StaticArg(axis)],
            axis=axis)

  Each harness can have a list of Limitations that describe the cases when
  the harness may not be fully implemented.


Encodes conditions under which a harness is limited, e.g., not runnable in JAX.

  See the module docstring for an introduction to harnesses and limitations.


Base class with helper functions for backward compatibility tests.


  Describes the dimension number arguments to an `XLA's Gather operator
  <https://www.openxla.org/xla/operation_semantics#gather>`_. See the XLA
  documentation for more details of what the dimension numbers mean.

  Args:
    offset_dims: the set of dimensions in the `gather` output that offset into
      an array sliced from `operand`. Must be a tuple of integers in ascending
      order, each representing a dimension number of the output.
    collapsed_slice_dims: the set of dimensions `i` in `operand` that have
      `slice_sizes[i] == 1` and that should not have a corresponding dimension
      in the output of the gather. Must be a tuple of integers in ascending
      order.
    start_index_map: for each dimension in `start_indices`, gives the
      corresponding dimension in the `operand` that is to be sliced. Must be a
      tuple of integers with size equal to `start_indices.shape[-1]`.
    operand_batching_dims: the set of batching dimensions `i` in `operand` that
      have `slice_sizes[i] == 1` and that should have a corresponding dimension
      in both the `start_indices` (at the same index in
      `start_indices_batching_dims`) and output of the gather. Must be a tuple
      of integers in ascending order.
    start_indices_batching_dims: the set of batching dimensions `i` in
      `start_indices` that should have a corresponding dimension in both the
      `operand` (at the same index in `operand_batching_dims`) and output of the
      gather. Must be a tuple of integers (order is fixed based on
      correspondence with `operand_batching_dims`).

  Unlike XLA's `GatherDimensionNumbers` structure, `index_vector_dim` is
  implicit; there is always an index vector dimension and it must always be the
  last dimension. To gather scalar indices, add a trailing dimension of size 1.


  Describes how to handle out-of-bounds indices in a gather or scatter.

  Possible values are:

  CLIP:
    Indices will be clamped to the nearest in-range value, i.e., such that the
    entire window to be gathered is in-range.
  FILL_OR_DROP:
    If any part of a gathered window is out of bounds, the entire window
    that is returned, even those elements that were otherwise in-bounds, will be
    filled with a constant.
    If any part of a scattered window is out of bounds, the entire window
    will be discarded.
  PROMISE_IN_BOUNDS:
    The user promises that indices are in bounds. No additional checking will be
    performed. In practice, with the current XLA implementation this means
    that out-of-bounds gathers will be clamped but out-of-bounds scatters will
    be discarded. Gradients will not be correct if indices are out-of-bounds.


  Describes the dimension number arguments to an `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_. See the XLA
  documentation for more details of what the dimension numbers mean.

  Args:
    update_window_dims: the set of dimensions in the `updates` that are window
      dimensions. Must be a tuple of integers in ascending
      order, each representing a dimension number.
    inserted_window_dims: the set of size 1 window dimensions that must be
      inserted into the shape of `updates`. Must be a tuple of integers in
      ascending order, each representing a dimension number of the output. These
      are the mirror image of `collapsed_slice_dims` in the case of `gather`.
    scatter_dims_to_operand_dims: for each dimension in `scatter_indices`, gives
      the corresponding dimension in `operand`. Must be a sequence of integers
      with size equal to `scatter_indices.shape[-1]`.
    operand_batching_dims: the set of batching dimensions `i` in `operand` that
      should have a corresponding dimension in both the `scatter_indices` (at
      the same index in `scatter_indices_batching_dims`) and `updates`. Must be
      a tuple of integers in ascending order. These are the mirror image of
      `operand_batching_dims` in the case of `gather`.
    scatter_indices_batching_dims: the set of batching dimensions `i` in
      `scatter_indices` that should have a corresponding dimension in both the
      `operand` (at the same index in `operand_batching_dims`) and output of the
      gather. Must be a tuple of integers (order is fixed based on
      correspondence with `input_batching_dims`). These are the mirror image of
      `start_indices_batching_dims` in the case of `gather`.

  Unlike XLA's `ScatterDimensionNumbers` structure, `index_vector_dim` is
  implicit; there is always an index vector dimension and it must always be the
  last dimension. To scatter scalar indices, add a trailing dimension of size 1.


Describes batch, spatial, and feature dimensions of a convolution.

  Args:
    lhs_spec: a tuple of nonnegative integer dimension numbers containing
      `(batch dimension, feature dimension, spatial dimensions...)`.
    rhs_spec: a tuple of nonnegative integer dimension numbers containing
      `(out feature dimension, in feature dimension, spatial dimensions...)`.
    out_spec: a tuple of nonnegative integer dimension numbers containing
      `(batch dimension, feature dimension, spatial dimensions...)`.


Rounding strategies for handling halfway values (e.g., 0.5) in
  :func:`jax.lax.round`.


Specify the tolerances used for computing unary functions.

  Maximum two tolerances can be specified: (atol and rtol) or (atol and ulps).