Precision enum for lax matrix multiply related functions.

  The device-dependent `precision` argument to JAX functions generally
  controls the tradeoff between speed and accuracy for array computations on
  accelerator backends, (i.e. TPU and GPU). Has no impact on CPU backends.
  This only has an effect on float32 computations, and does not affect the
  input/output datatypes. Members are:

  DEFAULT:
    Fastest mode, but least accurate. On TPU: performs float32 computations in
    bfloat16. On GPU: uses tensorfloat32 if available (e.g. on A100 and H100
    GPUs), otherwise standard float32 (e.g. on V100 GPUs). Aliases:
    ``'default'``, ``'fastest'``.
  HIGH:
    Slower but more accurate. On TPU: performs float32 computations in 3
    bfloat16 passes. On GPU: uses tensorfloat32 where available, otherwise
    float32. Aliases: ``'high'``..
  HIGHEST:
    Slowest but most accurate. On TPU: performs float32 computations in 6
    bfloat16. Aliases: ``'highest'``. On GPU: uses float32.


Specify the algorithm used for computing dot products.

  When used to specify the ``precision`` input to :func:`~jax.lax.dot`,
  :func:`~jax.lax.dot_general`, and other dot product functions, this data
  structure is used for controlling the properties of the algorithm used for
  computing the dot product. This API controls the precision used for the
  computation, and allows users to access hardware-specific accelerations.

  Support for these algorithms is platform dependent, and using an unsupported
  algorithm will raise a Python exception when the computation is compiled. The
  algorithms that are known to be supported on at least some platforms are
  listed in the :class:`~jax.lax.DotAlgorithmPreset` enum, and these are a
  good starting point for experimenting with this API.

  A "dot algorithm" is specified by the following parameters:

  * ``lhs_precision_type`` and ``rhs_precision_type``, the data types that the
    LHS and RHS of the operation are rounded to.
  * ``accumulation_type`` the data type used for accumulation.
  * ``lhs_component_count``, ``rhs_component_count``, and
    ``num_primitive_operations`` apply to algorithms that decompose the LHS
    and/or RHS into multiple components and execute multiple operations on
    those values, usually to emulate a higher precision. For algorithms with no
    decomposition, these values should be set to ``1``.
  * ``allow_imprecise_accumulation`` to specify if accumulation in lower
    precision is permitted for some steps (e.g.
    ``CUBLASLT_MATMUL_DESC_FAST_ACCUM``).

  The `StableHLO spec <https://openxla.org/stablehlo/spec#dot_general>`_ for
  the dot operation doesn't require that the precision types be the same as the
  storage types for the inputs or outputs, but some platforms may require that
  these types match. Furthermore, the return type of
  :func:`~jax.lax.dot_general` is always defined by the ``accumulation_type``
  parameter of the input algorithm, if specified.

  Examples:

    Accumulate two 16-bit floats using a 32-bit float accumulator:

    >>> algorithm = DotAlgorithm(
    ...     lhs_precision_type=np.float16,
    ...     rhs_precision_type=np.float16,
    ...     accumulation_type=np.float32,
    ... )
    >>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
    >>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
    >>> dot(lhs, rhs, precision=algorithm)  # doctest: +SKIP
    array([ 1.,  4.,  9., 16.], dtype=float16)

    Or, equivalently, using a preset:

    >>> algorithm = DotAlgorithmPreset.F16_F16_F32
    >>> dot(lhs, rhs, precision=algorithm)  # doctest: +SKIP
    array([ 1.,  4.,  9., 16.], dtype=float16)

    Presets can also be specified by name:

    >>> dot(lhs, rhs, precision="F16_F16_F32")  # doctest: +SKIP
    array([ 1.,  4.,  9., 16.], dtype=float16)

    The ``preferred_element_type`` parameter can be used to return the output
    without downcasting the accumulation type:

    >>> dot(lhs, rhs, precision="F16_F16_F32", preferred_element_type=np.float32)  # doctest: +SKIP
    array([ 1.,  4.,  9., 16.], dtype=float32)


An enum of known algorithms for computing dot products.

  This ``Enum`` provides a named set of :class:`~jax.lax.DotAlgorithm` objects
  that are known to be supported on at least platform. See the
  :class:`~jax.lax.DotAlgorithm` documentation for more details about the
  behavior of these algorithms.

  An algorithm can be selected from this list when calling :func:`~jax.lax.dot`,
  :func:`~jax.lax.dot_general`, or most other JAX dot product functions, by
  passing either a member of this ``Enum`` or it's name as a string using the
  ``precision`` argument.

  For example, users can specify the preset using this ``Enum`` directly:

  >>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
  >>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
  >>> algorithm = DotAlgorithmPreset.F16_F16_F32
  >>> dot(lhs, rhs, precision=algorithm)  # doctest: +SKIP
  array([ 1.,  4.,  9., 16.], dtype=float16)

  or, equivalently, they can be specified by name:

  >>> dot(lhs, rhs, precision="F16_F16_F32")  # doctest: +SKIP
  array([ 1.,  4.,  9., 16.], dtype=float16)

  The names of the presets are typically ``LHS_RHS_ACCUM`` where ``LHS`` and
  ``RHS`` are the element types of the ``lhs`` and ``rhs`` inputs
  respectively, and ``ACCUM`` is the element type of the accumulator. Some
  presets have an extra suffix, and the meaning of each of these is
  documented below. The supported presets are:


Describes ragged, group, and dot dimensions for ragged dot general.

  Args:
    dot_dimension_numbers: a tuple of tuples of sequences of ints of the form
      `((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims,
      rhs_batch_dims))`.
    lhs_ragged_dimensions: a sequence of ints indicating the 'lhs' ragged
      dimensions.
    rhs_group_dimensions: a sequence of ints indicating the 'rhs' group
      dimensions.


Describes which PRNG algorithm to use for rng_bit_generator.


Enum for SVD algorithm.


Describes which FFT operation to perform.


A bounded functional stack implementation. Elements may be pytrees.


Describes a subproblem of _eigh_work.

  Each subproblem is a `size` x `size` Hermitian matrix, starting at `offset`
  in the workspace.


Interpolate points on a regular rectangular grid.

  JAX implementation of :func:`scipy.interpolate.RegularGridInterpolator`.

  Args:
    points: length-N sequence of arrays specifying the grid coordinates.
    values: N-dimensional array specifying the grid values.
    method: interpolation method, either ``"linear"`` or ``"nearest"``.
    bounds_error: not implemented by JAX
    fill_value: value returned for points outside the grid, defaults to NaN.

  Returns:
    interpolator: callable interpolation object.

  Examples:
    >>> points = (jnp.array([1, 2, 3]), jnp.array([4, 5, 6]))
    >>> values = jnp.array([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
    >>> interpolate = RegularGridInterpolator(points, values, method='linear')

    >>> query_points = jnp.array([[1.5, 4.5], [2.2, 5.8]])
    >>> interpolate(query_points)
    Array([30., 64.], dtype=float32)


Universal functions which operation element-by-element on arrays.

  JAX implementation of :class:`numpy.ufunc`.

  This is a class for JAX-backed implementations of NumPy's ufunc APIs.
  Most users will never need to instantiate :class:`ufunc`, but rather
  will use the pre-defined ufuncs in :mod:`jax.numpy`.

  For constructing your own ufuncs, see :func:`jax.numpy.frompyfunc`.

  Examples:
    Universal functions are functions that apply element-wise to broadcasted
    arrays, but they also come with a number of extra attributes and methods.

    As an example, consider the function :obj:`jax.numpy.add`. The object
    acts as a function that applies addition to broadcasted arrays in an
    element-wise manner:

    >>> x = jnp.array([1, 2, 3, 4, 5])
    >>> jnp.add(x, 1)
    Array([2, 3, 4, 5, 6], dtype=int32)

    Each :class:`ufunc` object includes a number of attributes that describe
    its behavior:

    >>> jnp.add.nin  # number of inputs
    2
    >>> jnp.add.nout  # number of outputs
    1
    >>> jnp.add.identity  # identity value, or None if no identity exists
    0

    Binary ufuncs like :obj:`jax.numpy.add` include  number of methods to
    apply the function to arrays in different manners.

    The :meth:`~ufunc.outer` method applies the function to the
    pair-wise outer-product of the input array values:

    >>> jnp.add.outer(x, x)
    Array([[ 2,  3,  4,  5,  6],
           [ 3,  4,  5,  6,  7],
           [ 4,  5,  6,  7,  8],
           [ 5,  6,  7,  8,  9],
           [ 6,  7,  8,  9, 10]], dtype=int32)

    The :meth:`ufunc.reduce` method performs a reduction over the array.
    For example, :meth:`jnp.add.reduce` is equivalent to ``jnp.sum``:

    >>> jnp.add.reduce(x)
    Array(15, dtype=int32)

    The :meth:`ufunc.accumulate` method performs a cumulative reduction
    over the array. For example, :meth:`jnp.add.accumulate` is equivalent
    to :func:`jax.numpy.cumulative_sum`:

    >>> jnp.add.accumulate(x)
    Array([ 1,  3,  6, 10, 15], dtype=int32)

    The :meth:`ufunc.at` method applies the function at particular indices in the
    array; for ``jnp.add`` the computation is similar to :func:`jax.lax.scatter_add`:

    >>> jnp.add.at(x, 0, 100, inplace=False)
    Array([101,   2,   3,   4,   5], dtype=int32)

    And the :meth:`ufunc.reduceat` method performs a number of ``reduce``
    operations between specified indices of an array; for ``jnp.add`` the
    operation is similar to :func:`jax.ops.segment_sum`:

    >>> jnp.add.reduceat(x, jnp.array([0, 2]))
    Array([ 3, 12], dtype=int32)

    In this case, the first element is ``x[0:2].sum()``, and the second element
    is ``x[2:].sum()``.


Return dense multi-dimensional "meshgrid".

  LAX-backend implementation of :obj:`numpy.mgrid`. This is a convenience wrapper for
  functionality provided by :func:`jax.numpy.meshgrid` with ``sparse=False``.

  See Also:
    jnp.ogrid: open/sparse version of jnp.mgrid

  Examples:
    Pass ``[start:stop:step]`` to generate values similar to :func:`jax.numpy.arange`:

    >>> jnp.mgrid[0:4:1]
    Array([0, 1, 2, 3], dtype=int32)

    Passing an imaginary step generates values similar to :func:`jax.numpy.linspace`:

    >>> jnp.mgrid[0:1:4j]
    Array([0.        , 0.33333334, 0.6666667 , 1.        ], dtype=float32)

    Multiple slices can be used to create broadcasted grids of indices:

    >>> jnp.mgrid[:2, :3]
    Array([[[0, 0, 0],
            [1, 1, 1]],
           [[0, 1, 2],
            [0, 1, 2]]], dtype=int32)


Return open multi-dimensional "meshgrid".

  LAX-backend implementation of :obj:`numpy.ogrid`. This is a convenience wrapper for
  functionality provided by :func:`jax.numpy.meshgrid` with ``sparse=True``.

  See Also:
    jnp.mgrid: dense version of jnp.ogrid

  Examples:
    Pass ``[start:stop:step]`` to generate values similar to :func:`jax.numpy.arange`:

    >>> jnp.ogrid[0:4:1]
    Array([0, 1, 2, 3], dtype=int32)

    Passing an imaginary step generates values similar to :func:`jax.numpy.linspace`:

    >>> jnp.ogrid[0:1:4j]
    Array([0.        , 0.33333334, 0.6666667 , 1.        ], dtype=float32)

    Multiple slices can be used to create sparse grids of indices:

    >>> jnp.ogrid[:2, :3]
    [Array([[0],
            [1]], dtype=int32),
     Array([[0, 1, 2]], dtype=int32)]


Concatenates slices, scalars and array-like objects along a given axis.


Concatenate slices, scalars and array-like objects along the first axis.

  LAX-backend implementation of :obj:`numpy.r_`.

  See Also:
    ``jnp.c_``: Concatenates slices, scalars and array-like objects along the last axis.

  Examples:
    Passing slices in the form ``[start:stop:step]`` generates ``jnp.arange`` objects:

    >>> jnp.r_[-1:5:1, 0, 0, jnp.array([1,2,3])]
    Array([-1,  0,  1,  2,  3,  4,  0,  0,  1,  2,  3], dtype=int32)

    An imaginary value for ``step`` will create a ``jnp.linspace`` object instead,
    which includes the right endpoint:

    >>> jnp.r_[-1:1:6j, 0, jnp.array([1,2,3])]
    Array([-1.        , -0.6       , -0.20000002,  0.20000005,
           0.6       ,  1.        ,  0.        ,  1.        ,
           2.        ,  3.        ], dtype=float32)

    Use a string directive of the form ``"axis,dims,trans1d"`` as the first argument to
    specify concatenation axis, minimum number of dimensions, and the position of the
    upgraded array's original dimensions in the resulting array's shape tuple:

    >>> jnp.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, 2D output
    Array([[1, 2, 3],
           [4, 5, 6]], dtype=int32)

    >>> jnp.r_['0,2,0', [1,2,3], [4,5,6]] # push last input axis to the front
    Array([[1],
           [2],
           [3],
           [4],
           [5],
           [6]], dtype=int32)

    Negative values for ``trans1d`` offset the last axis towards the start
    of the shape tuple:

    >>> jnp.r_['0,2,-2', [1,2,3], [4,5,6]]
    Array([[1],
           [2],
           [3],
           [4],
           [5],
           [6]], dtype=int32)

    Use the special directives ``"r"`` or ``"c"`` as the first argument on flat inputs
    to create an array with an extra row or column axis, respectively:

    >>> jnp.r_['r',[1,2,3], [4,5,6]]
    Array([[1, 2, 3, 4, 5, 6]], dtype=int32)

    >>> jnp.r_['c',[1,2,3], [4,5,6]]
    Array([[1],
           [2],
           [3],
           [4],
           [5],
           [6]], dtype=int32)

    For higher-dimensional inputs (``dim >= 2``), both directives ``"r"`` and ``"c"``
    give the same result.


Concatenate slices, scalars and array-like objects along the last axis.

  LAX-backend implementation of :obj:`numpy.c_`.

  See Also:
    ``jnp.r_``: Concatenates slices, scalars and array-like objects along the first axis.

  Examples:

    >>> a = jnp.arange(6).reshape((2,3))
    >>> jnp.c_[a,a]
    Array([[0, 1, 2, 0, 1, 2],
           [3, 4, 5, 3, 4, 5]], dtype=int32)

    Use a string directive of the form ``"axis:dims:trans1d"`` as the first argument to specify
    concatenation axis, minimum number of dimensions, and the position of the upgraded array's
    original dimensions in the resulting array's shape tuple:

    >>> jnp.c_['0,2', [1,2,3], [4,5,6]]
    Array([[1],
           [2],
           [3],
           [4],
           [5],
           [6]], dtype=int32)

    >>> jnp.c_['0,2,-1', [1,2,3], [4,5,6]]
    Array([[1, 2, 3],
           [4, 5, 6]], dtype=int32)

    Use the special directives ``"r"`` or ``"c"`` as the first argument on flat inputs
    to create an array with inputs stacked along the last axis:

    >>> jnp.c_['r',[1,2,3], [4,5,6]]
    Array([[1, 4],
           [2, 5],
           [3, 6]], dtype=int32)


Struct returned by :func:`jax.numpy.unique_all`.


A context manager to set the error checking behavior.

  If both `all` and a category are provided, the category will override the
  `all` setting.

  When the error checking behavior is set to "ignore", all errors will be
  ignored. When set to "raise", errors will be detected and recorded, but an
  exception will not be raised immediately. Users must call
  :func:`raise_if_error` to at the end of the computation to raise the
  exception.


Helper property for index update functionality.

  The ``at`` property provides a functionally pure equivalent of in-place
  array modifications.

  In particular:

  ==============================  ================================
  Alternate syntax                Equivalent In-place expression
  ==============================  ================================
  ``x = x.at[idx].set(y)``        ``x[idx] = y``
  ``x = x.at[idx].add(y)``        ``x[idx] += y``
  ``x = x.at[idx].subtract(y)``   ``x[idx] -= y``
  ``x = x.at[idx].multiply(y)``   ``x[idx] *= y``
  ``x = x.at[idx].divide(y)``     ``x[idx] /= y``
  ``x = x.at[idx].power(y)``      ``x[idx] **= y``
  ``x = x.at[idx].min(y)``        ``x[idx] = minimum(x[idx], y)``
  ``x = x.at[idx].max(y)``        ``x[idx] = maximum(x[idx], y)``
  ``x = x.at[idx].apply(ufunc)``  ``ufunc.at(x, idx)``
  ``x = x.at[idx].get()``         ``x = x[idx]``
  ==============================  ================================

  None of the ``x.at`` expressions modify the original ``x``; instead they return
  a modified copy of ``x``. However, inside a :py:func:`~jax.jit` compiled function,
  expressions like :code:`x = x.at[idx].set(y)` are guaranteed to be applied in-place.

  Unlike NumPy in-place operations such as :code:`x[idx] += y`, if multiple
  indices refer to the same location, all updates will be applied (NumPy would
  only apply the last update, rather than applying all updates.) The order
  in which conflicting updates are applied is implementation-defined and may be
  nondeterministic (e.g., due to concurrency on some hardware platforms).

  By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound
  index semantics can be specified via the ``mode`` parameter (see below).

  Args:
    mode: string specifying out-of-bound indexing mode. Options are:

      - ``"promise_in_bounds"``: (default) The user promises that indices are in bounds.
        No additional checking will be performed. In practice, this means that
        out-of-bounds indices in ``get()`` will be clipped, and out-of-bounds indices
        in ``set()``, ``add()``, etc. will be dropped.
      - ``"clip"``: clamp out of bounds indices into valid range.
      - ``"drop"``: ignore out-of-bound indices.
      - ``"fill"``: alias for ``"drop"``.  For `get()`, the optional ``fill_value``
        argument specifies the value that will be returned.

      See :class:`jax.lax.GatherScatterMode` for more details.
    wrap_negative_indices: If True (default) then negative indices indicate position
      from the end of the array, similar to Python and NumPy indexing. If False, then
      negative indices are considered out-of-bounds and behave according to the
      ``mode`` parameter.
    fill_value: Only applies to the ``get()`` method: the fill value to return for
      out-of-bounds slices when ``mode`` is ``'fill'``. Ignored otherwise. Defaults
      to ``NaN`` for inexact types, the largest negative value for signed types, the
      largest positive value for unsigned types, and ``True`` for booleans.
    indices_are_sorted: If True, the implementation will assume that the (normalized)
      indices passed to ``at[]`` are sorted in ascending order, which can lead to more
      efficient execution on some backends. If True but the indices are not actually
      sorted, the output is undefined.
    unique_indices: If True, the implementation will assume that the (normalized) indices
      passed to ``at[]`` are unique, which can result in more efficient execution on some
      backends. If True but the indices are not actually unique, the output is undefined.

  Examples:
    >>> x = jnp.arange(5.0)
    >>> x
    Array([0., 1., 2., 3., 4.], dtype=float32)
    >>> x.at[2].get()
    Array(2., dtype=float32)
    >>> x.at[2].add(10)
    Array([ 0.,  1., 12.,  3.,  4.], dtype=float32)

    By default, out-of-bound indices are ignored in updates, but this behavior
    can be controlled with the ``mode`` parameter:

    >>> x.at[10].add(10)  # dropped
    Array([0., 1., 2., 3., 4.], dtype=float32)
    >>> x.at[20].add(10, mode='clip')  # clipped
    Array([ 0.,  1.,  2.,  3., 14.], dtype=float32)

    For ``get()``, out-of-bound indices are clipped by default:

    >>> x.at[20].get()  # out-of-bounds indices clipped
    Array(4., dtype=float32)
    >>> x.at[20].get(mode='fill')  # out-of-bounds indices filled with NaN
    Array(nan, dtype=float32)
    >>> x.at[20].get(mode='fill', fill_value=-1)  # custom fill value
    Array(-1., dtype=float32)

    Negative indices count from the end of the array, but this behavior can
    be disabled by setting ``wrap_negative_indices = False``:

    >>> x.at[-1].set(99)
    Array([ 0.,  1.,  2.,  3., 99.], dtype=float32)
    >>> x.at[-1].set(99, wrap_negative_indices=False, mode='drop')  # dropped!
    Array([0., 1., 2., 3., 4.], dtype=float32)


Helper object to call indexed update functions for an (advanced) index.

  This object references a source array and a specific indexer into that array.
  Methods on this object return copies of the source array that have been
  modified at the positions specified by the indexer.


Metadata for the `Python array API`_

  .. _Python array API: https://data-apis.org/array-api/


Common dtype for all kinds of semaphore dtypes.

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


Regular semaphore dtype.

  Like its superclass, this class should never be instantiated.


Barrier semaphore dtype.

  Like its superclass, this class should never be instantiated.


Base class for compiler parameters.


Specifies how a block should be buffered for a pipeline.

  Attributes:
    buffer_count: The number of buffers to use for multiple buffering.
    use_lookahead: optional bool, indicates whether to use lookahead on the
      buffer. Enabling lookahead allows the pipeline to begin fetching the next
      changed block as soon as a slot is available, no matter how many
      iterations ahead that block is.


Like jax.ShapeDtypeStruct but with memory spaces.


Logical, device-agnostic memory spaces.

  Each memory space will be translated to a device-specific memory
  type during lowering.


Use to index an array using an elementwise start index.


The default BlockShape type.


Allows to specify a bounded slice of a dimension.

  Specifically, the index_map need to return a `pl.Slice/pl.ds` for this
  dimension. The start and size may be dynamic, as long as the size <=
  block_size.


Helper class that checks for index_map equality.


Specifies how an array should be sliced for each invocation of a kernel.

  The `block_shape` is a sequence of `int | None`s, or `BlockDim` types (e.g.
  `pl.Element`, `pl.Squeezed`, `pl.Blocked`, `pl.BoundedSlice`). Each of these
  types specify the size of the block dimension. `None` is used to specify a
  dimension that is squeezed out of the kernel. The `BlockDim` types allow for
  more fine-grained control over the indexing of the dimension. The `index_map`
  needs to return a tuple of the same length as `block_shape`, which each entry
  depending on the type of `BlockDim`.

  See :ref:`pallas_blockspec` and the individual `BlockDim` type docstrings for
  more details.


Transforms a memory reference on load or store.


An internal canonicalized version of BlockSpec.

  See the `check_invariants` method for precise specification.


An internal canonicalized version of GridSpec.

  Encodes the calling conventions of the pallas_call primitive, the kernel,
  and the index maps.

  The pallas_call is invoked with: ``*dynamic_grid_sizes, *index, *inputs``.
  The ``index`` operands are for the scalar prefetch.

  The kernel function is invoked with:
  ``*index, *inputs, *scratch``.

  The index map functions are invoked with:
  ``*program_ids, *index``.

  See the `check_invariants` method for a more precise specification.


Encodes the grid parameters for :func:`jax.experimental.pallas.pallas_call`.

  See the documentation for :func:`jax.experimental.pallas.pallas_call`,
  and also :ref:`pallas_grids_and_blockspecs` for a more detailed
  description of the parameters.


A model of a memory reference.

  When a reference has a small leading dimension, it might be represented by
  multiple slots in the reference array. Its region starts at base (that can be
  dynamic) and has the given length (always static).


A model of a semaphore reference.

  Semaphore arrays are always fully unrolled and are represented by a contiguous
  subset of the global semaphore array.


Error thrown by Pallas when re-raising a verification error.


Mosaic TPU compiler parameters.

  Attributes:
    dimension_semantics: A list of dimension semantics for each grid dimension
      of the kernel. Either "parallel" for dimensions that can execute in any
      order, or "arbitrary" for dimensions that must be executed sequentially.
    allow_input_fusion: A list of booleans indicating whether input fusion is
      allowed for each argument.
    vmem_limit_bytes: Overrides the default VMEM limit for a kernel. Note that
      this must be used in conjunction with the
      --xla_tpu_scoped_vmem_limit_kib=N flag with N*1kib > vmem_limit_bytes.
    collective_id: Indicates which barrier semaphore to use for the kernel. Note
      that using the same collective_id does not guarantee that the same barrier
      semaphore will be allocated between kernels.
    has_side_effects: Set to True to prevent kernel being CSEd by XLA.
    flags: A dictionary of command line flags for the kernel.
    internal_scratch_in_bytes: The size of the internal scratch space used by
      Mosaic.
    serialization_format: The serialization format for the kernel body.
    kernel_type: Specify if the kernel is meant to run on TensorCore or one of
      the SparseCores
    disable_bounds_checks: Disable bounds checks in the kernel.


A mesh of TensorCores.


Buffer type for the arguments to an emitted pipeline.


Abstract interface for BufferedRefs.


A helper class to automate VMEM double buffering in pallas pipelines.

  Attributes:
    spec: pallas blockspec.
    dtype: dtype for buffers.
    buffer_type: enum indicating whether this is an input, output, or in/out
      accumulator buffered reference.
    window_ref: a multiple-buffer to hold the working and dirty buffers used
      to copy into and out of.  In the case of a BufferedRef targeting a VMEM
      reference, this simply points to the existing ref.
    accum_ref: accumulating buffer used by accumulator BufferedRefs.
    copy_in_slot: current slot to copy in for the working buffer.
    copy_out_slot: current slot to copy out for the working buffer.
    wait_in_slot: current slot to wait in for the working buffer.
    wait_out_slot: current slot to wait out for the working buffer.
    next_fetch_smem: Holds the next grid indices to fetch for lookahead. This
      is the SMEM backing buffer used to persist state between pipeline
      invocations.
    next_fetch_sreg: Holds the next grid indices to fetch for lookahead. This
      is the register state used to track the indices within the pipeline loop.
    sem_recvs: Multiple buffered semaphores for input DMAs.
    sem_sends: Multiple buffered semaphores for output DMAs.
    block_shape: passthrough property for the BlockSpec's block_shape.
    compute_index: passthrough property for the BlockSpec's compute_index.
    memory_space: passthrough property for the BlockSpec's memory_space.
    current_ref: points to the current working slice of the double-buffer.
    is_input: whether this BufferedRef acts as a pipeline input.
    is_output: whether this BufferedRef acts as a pipeline output.
    is_accumulator: whether this BufferedRef is an accumulator.
    is_input_output: whether this BufferedRef is an input/output without
      automatic accumulation.
    swap: Tracks whether the BufferedRef slots need to be swapped before next
      copy.


Sequences input and output copies and waits for a pipeline.


Parameters for TPU interpret mode.

  TPU interpret mode is a way run Pallas TPU kernels on CPU, while simulating
  a TPU's shared memory (HBM, VMEM, etc.), communication (remote and local
  DMAs), and synchronization operations (semaphores, barriers, etc.).  This mode
  is intended for debugging and testing.

  To run a kernel under TPU interpret mode, pass an instance of
  ``InterpretParams`` as an argument for the ``interpret`` parameter of
  :func:`jax.experimental.pallas.pallas_call` or
  :func:`jax.experimental.pallas.core_map`.

  Attributes:
    dma_execution_mode:  If "eager", DMAs are executed as soon as they are
      issued.  If "on_wait", DMA reads or writes are only executed when a device
      is waiting on a DMA semaphore that will be signaled when the read or write
      is complete.
      Default: "on_wait".
    detect_races: If True, a dynamic, happens-before race detector will be used
      to detect data races during kernel interpretation.  If any races are
      detected, a message will be printed and `races.races_found` will be set to
      True.
      Default: False.
    out_of_bounds_reads: If "raise", an exception will be raised on any
      out-of-bounds read of a buffer.  If "uninitialized_value", any parts of
      the read that are out-of-bounds will return the value used to fill
      uninitialized memory, which can be configured via the
      "uninitialized_memory".  NOTE: If an exception is raised while
      interpreting a kernel, you must call
      :func:`reset_tpu_interpret_mode_state` before using TPU interpret mode
      again in the same process.
      Default: "raise".
    skip_floating_point_ops: If True, operations that produce only floating
      point values will not be interpreted; instead, their results will be
      replaced with arrays all of `jnp.inf`. Additionally any floating point
      operands to any operation will be replaced with (arrays of) `jnp.inf`.
      Default: False.
    uninitialized_memory: If "nan", allocated buffers are initialized to contain
      all NaNs (or to their maximum possible value for integers). If "zero",
      allocated buffers are initialized to all zeros.
      Default: "nan".
    random_seed: Seed for random number generator used during interpretation.
      Currently random numbers are used to randomize the grid coordinates along
      dimensions with 'parallel' semantics.
      Default: None.
    grid_point_recorder: Callback that is invoked by the interpreter for each
      grid point in the order in which the grid points are traversed. The
      callback is invoked with two arguments:
        - A tuple of grid coordinates.
        - The local core ID of the core that is processing the grid point.
      This callback is intended for inspecting
        - the randomization of coordinates along grid dimensions with 'parallel'
          semantics and
        - the mapping of grid points to local (i.e. per-device) cores.
      Default: None.
    num_cores_per_device: The number of cores per device.
      Default: 1.


A container class for PRNG key data.

  We pass around keys as a KeyScalarBundle in the lowering pass rather than
  as a vector, since we want the key data to live in scalar registers rather
  than vector registers. This special dataclass exists so we can return
  multiple scalar values from load_op, because the load_op primitive does
  not allow multiple results.

  Attributes:
    scalars: A list of OpResults representing scalar key data during the
      lowering pass.


Protocol for a compute context for the warp specialized pipeline.

  The ComputeContext is run exclusively in the compute thread and allows
  the user to set up a prologue to initialize a pipeline carry and an epilogue
  to consume the final carry.

  All values allocated in the ComputeContext will only be allocated in the
  compute thread and not the memory thread. This can potentially reduce
  register pressure if certain values are only consumed by the compute threads.

  Usage will usually follow this structure:

  ```
  def compute_context(pipeline):
    # Perform prologue work and compute the initial carry.
    initial_carry = ...
    # Run the pipeline.
    final_carry = pipeline(*initial_carry)
    # Perform epilogue work using the final carry.
    do_work(final_carry)
  ```


Mosaic GPU compiler parameters.

  Attributes:
    approx_math: If True, the compiler is allowed to use approximate
      implementations of some math operations, e.g. ``exp``. Defaults to False.
    dimension_semantics: A list of dimension semantics for each grid
      dimension of the kernel. Either "parallel" for dimensions that can
      execute in any order, or "sequential" for dimensions that must be
      executed sequentially.
    max_concurrent_steps: The maximum number of sequential stages that are
      active concurrently. Defaults to 1.
    delay_release: The number of steps to wait before reusing the input/output
      references. Defaults to 0, and must be strictly smaller than
      max_concurrent_steps. Generally, you'll want to set it to 1 if you don't
      await the WGMMA in the body.
    unsafe_no_auto_barriers: If True, Pallas will never automatically insert
      barrier instructions that ensure synchronous semantics of loads and stores.
      At the moment, the insertion is done conservatively and might regress
      performance. There are (at least) two conditions that must be satisfied
      for the use of this flag to be safe. First, no memory region is ever read
      *and* written to by the same thread (async copies are performed by
      background threads and do not count towards this rule). Secondly, no
      thread ever calls commit_smem(), reads from the committed SMEM and then
      issues an async copy overwriting that region (this is a very artificial
      and highly unlikely scenario).
    profile_space: The number of profiler events that can be collected in a
      single invocation. It is undefined behavior if a thread collects more
      events than this.
    profile_dir: The directory to which profiling traces will be written to.


Thread semantics for a primitives at the Pallas user-level.


A sequence of trees of refs that are allowed to reuse the same memory.

  One should not make assumptions as to how each ref will map to the underlying
  memory region, since arbitrary padding may be applied in between different
  refs.

  As such, ref unions are only safe to use when the groups of refs that we
  intend to alias have disjoint lifetimes (i.e. one should never attempt to read
  data using a different ref than the one that was used to write the data).


Represents a tiling transformation for memory refs.

  A tiling of (X, Y) on an array of shape (M, N) will result in a transformed
  shape of (M // X, N // Y, X, Y). Ex. A (256, 256) block that is tiled with a
  tiling of (64, 32) will be tiled as (4, 8, 64, 32).


Transpose a tiled memref.


Bitcasts the underlying ref at the given offset to the given shape and dtype.


Describes a barrier reference.

  Attributes:
    num_arrivals: The number of arrivals that will be recorded by this barrier.
    num_barriers: The number of barriers that will be created. Individual
      barriers can be accessed by indexing into the barrier Ref.
    orders_tensor_core: If False, a successfull wait from one thread does not
      guarantee that the TensorCore-related operations in other threads have
      completed. Similarly, when False any TensorCore operation in the waiting
      thread is allowed to begin before the wait succeeds.


Represents a mesh over individual warps within a warpgroup.

  When used in conjunction with `core_map`, the warp ID will be visible
  within the body of the wrapped scope by querying `lax.axis_index` with
  the specified axis name.


Layout for TMEM references.


Compiler parameters for Triton.

  Attributes:
    num_warps: The number of warps to use for the kernel. Each warp consists of
      32 threads.
    num_stages: The number of stages the compiler should use for software
      pipelining loops.


Keeps python objects alive.


Base class for fusible extended dtypes.


Represents a connection to exactly one peer.


Interface for checkpointing GDAs asynchronously.

  This class manages the state of an ongoing asynchronous checkpoint.

  For example, say a checkpoint happens on every step. If you checkpoint on
  step 1 and after some computation the model is on checkpoint 2. But step 1's
  checkpoint hasn't finished committing to the storage layer yet. So until that
  is finished, checkpoint for step 2 will need to be blocked. Maintaining a
  class allows to maintain that state.

  Examples:

  Below is a simplified training loop:

  ```
  # Call this at the start of your program.
  jax.distributed.initialize()

  manager = GlobalAsyncCheckpointManager()

  # Restore checkpoint if available or initialize the train_state from
  # init_fn().
  train_state = manager.deserialize(...)

  while ...:
    if step % num_steps_between_checkpoints == 0:
      manager.serialize(train_state, temp_checkpoint_dir=...,
                        final_checkpoint_dir=...)
      train_state = train_step(train_state, input)
      # This is a non-blocking call.
      manager.check_for_errors()

  manager.serialize(train_state, temp_checkpoint_dir=...,
                    final_checkpoint_dir=...)
  # Wait before the end of the program for the checkpoint to finish. This is a
  # blocking call.
  manager.wait_until_finished()
  ```


Responsible for serializing GDAs via TensorStore.


Limits host scratch memory usage when reading/writing checkpoints per process.


A wrapper around a Future that makes it look like an async function.


A container for a source map and the paired generated code.


User function wrapped by colocated_python.


Specialization for a colocated_python function.


Registry of object instances.


Temporarily stores results from synchronous execution of functions.


Stores live objects.


SegmentIds for Q and KV sequences.

  SegmentIds are used to generate segment mask, which prevents attention between
  different segments in the input sequence. Each array is a list of ids
  (integers).
  Only the token with the same id can attend to each other.

  Attributes:
    q: segment ids along the Q sequence.
    kv: segment ids along the KV sequence.


Tile sizes parameterizing FlashAttention kernels.

  Those parameters have negligible effect on numerics, but affect performance
  greatly.


Contains runtime masking information for the Splash attention kernel.

  The arrays data_next, mask_next and block_mask are placed in TPU
  scalar-memory. This is a scarse resource so the mask creation logic attempts
  to shrink the data-type of these arrays to the smallest possible one.
  This can be: np.int32, np.int16 or np.int8.

  For the arrays data_next, mask_next and block_mask the size of the first
  dimension can be one of the two following values: num_head or
  num_head_shards.
  The first dimension has size:
  * num_head_shards when there is only one unique mask for each head in a shard.
  In this case the three arrays are broadcasted to all the heads in the shard.
  * num_heads when there is more than one unique mask for each head in the
  shard.

  Attributes:
    data_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
      NumPy array where each entry contains the next `kv` block index to
      prefetch.
    mask_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
      NumPy array where each entry contains the next mask block index in
      `partial_mask_blocks` to prefetch.
    block_mask: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
      NumPy array whose entries can be 0, 1 or 2. An entry of 0 indicates that
      the corresponding block in the full mask was all zeros. An entry of 1
      indicates that the corresponding block in the full mask contained both
      zeros and ones. An entry of 2 indicates the corresponding block was
      entirely ones.
    partial_mask_blocks: A bool[num_partial_blocks, block_q, block_kv] NumPy
      array that contains the blocks of the original mask that contained both
      zeros and ones. The entries in `mask_next` point to indices in the first
      axis of this array.
    q_sequence: A i32[q_sequence_length] NumPy array. When using causal masking,
      this contains the list of indices that correspond to q tokens. For plain
      causal this is just np.arange(q_sequence_length).
    is_dynamic_mask: A bool indicating whether the mask is dynamic or static.
      When True, the leading dimensions of `partial_mask_blocks` (num_heads,
      q_blocks, kv_blocks) are not collapsed, allowing us to shard it along
      those dimensions.


Helper to make a numpy array hashable: can be added associative containers.

  Attributes:
    array: The underlying numpy array.


SegmentIds for Q and KV sequences.

  SegmentIds are a mechanism to ensure that there is no cross-attention between
  segments (fraction of a sequence) that have been concatenated together into a
  sequence. Each array is a list of ids (integers). Only tokens with the same
  id are allowed to attend to each other.

  The static mask (e.g. causal) is "and-ed" with the segment id mask to form
  the actual attention mask. It is important that the latter does not have any
  all-zero rows (along dimension kv). Otherwise it would result in a invalid
  softmax (the denominator would be 0).
  This condition holds for causal self-attention because in this case segment
  ids form a block diagonal matrix so at least one element in each row is set.
  It is easy to break this condition with non-self-attention configurations.
  Attributes:
    q: segment ids along the Q sequence
    kv: segment ids along the KV sequence


Tile sizes parameterizing SplashAttention kernels.

  Those parameters have negligible effect on numerics, but affect performance
  greatly.

  Note that changing the layouts only influences the physical layout that the
  kernel will enforce. The logical interface to splash attention always takes
  the head dimension as the minormost one.


A base class for splash attention masks.


Lazy multihead mask, combines multiple lazy masks one per head.


Superclass for all masks that can be computed inside the kernel using a callable object.

  This subclass is designed to be used with Splash Attention.
  It allows the mask logic to be computed on-the-fly or fused into the attention
  kernel, avoiding the memory cost of materializing the full
  (sequence_length, sequence_length) boolean mask array, which can be excessive
  for long sequences.

  Attributes:
    _shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
    offset: Offset of q start wrt kv. A positive offset shifts the bottom
      triangle upward, a negative one shifts it downward. A negative offset
      makes the first 'offset' rows of the attention matrix all 0s which leads
      to undefined softmax.
    q_sequence: Indices of Q sequence. q_sequence is reused across __getitem__
      calls which is important for compile-time performance.
    mask_function: Function used by the SplashAttention kernel to compute the
      mask rather than loading it.


Lazy causal mask, prevents the model from attending to future tokens.

  Attributes:
    offset: Offset of q start wrt kv. A positive offset shifts the bottom
      triangle upward, a negative one shifts it downward. A negative offset
      makes the first 'offset' rows of the attention matrix all 0s which leads
      to undefined softmax.


Lazy chunked causal mask.

  Attention is causal within each chunk (0, K), (K, 2K), (2K, 3K), ... tokens
  attend to each other but not across chunks.
  Llama4 models use interleaved chunk attention along with global attention.

  Attributes:
    chunk_size: The size of each attention chunk.


Lazy local mask, prevents model from attending to tokens outside window.

  Attributes:
    window_size: Size of the two sides of the local window (None identifies no
      limit for the given side).
    offset: Offset of q start wrt kv. A positive offset shifts the bottom
      triangle upward, a negative one shifts it downward. A negative offset
      makes the first 'offset' rows of the attention matrix all 0s which leads
      to undefined softmax.


A mask backed by a dense numpy array.


Lazy full mask, allows all tokens to attend to all other tokens.


A tensor which has been quantized to int8 and its scales.

  Attributes:
    weight: Weight
    scales: Scales


Descriptor for async copy of multiple K/V pages from HBM.


  Tile sizes parameterizing the attention kernel. These block sizes
  should be tuned for the model and hardware for optimal performance.

  Attributes:
    block_q: Block size along Q sequence length for forward kernel.
    block_k: Block size along KV sequence length for forward kernel.
    block_kv: Block size along KV sequence length for forward kernel.
    block_q_dkv: Block size along Q sequence length for dKV backward kernel.
    block_kv_dkv: Block size along KV sequence length for dKV backward kernel.
    block_q_dq: Block size along Q sequence length for dQ backward kernel.
    block_kv_dq: Block size along KV sequence length for dQ backward kernel.


Information regarding the group being processed in a block.


Tiles a suffix of memref dimensions.

  For example, given a memref of shape (5, 128, 128) and a tiling of (64, 32),
  the shape of the result will be (5, 2, 4, 64, 32). The shape always ends with
  the tile shape, and the size of tiled dimensions is divided by the tile size.
  This is especially useful for swizzled WGMMA, which expect tiled layouts in
  shared memory.


Transposes memref dimensions.


Collapses leading indices into one.


Manages ops handling the GMEM scratch that contains the TMA descriptors.

  TMA descriptors are created on the host and then copied to GMEM. So there
  needs to be some code on the host to allocate and initialize the TMA
  descriptors. However, we only know what descriptors we need after we have
  lowered the entire kernel. This class helps manage everything needed to
  correctly allocate and initialize the scratch.

  To help reconcile the needs of kernels that use the dialect lowering with
  those that use MGPU APIs directly, this class only creates the relevant ops
  lazily. Eager creation would make them appear dead before dialect lowering
  and MLIR's DCE would remove them.

  During the lowering, we collect information about how many bytes are needed
  and also how each descriptor should be initialized on the host. At the end
  of the lowering, the finalize_size() method should be called to add the
  necessary code on the host to allocate and initialize all descriptors.

  Here's how the IR looks after the initial ops are created for the first time:

  %1 = llvm.alloc_op {elem_type = !llvm.array<0 x i8>} -> !llvm.ptr
  %2 = llvm.load_op (%1) : (!llvm.ptr) -> !llvm.array<0 x i8>
  ...
  %3 = gpu.launch async
    ^bb0:
      %4 = builtin.unrealized_conversion_cast_op(%2)
             : (!llvm.array<256 x i8>) -> !llvm.ptr

  And here is an example of how the IR could look like after finalize_size() is
  called:

  %11 = llvm.alloc_op {elem_type = !llvm.array<256 x i8>} -> !llvm.ptr
  %22 = llvm.load_op (%11) : (!llvm.ptr) -> !llvm.array<256 x i8>
  ...
  # Ops inserted to initialize the tma descriptors on the host:
  ...
  %33 = llvm.getelementptr %11[0] : (!llvm.ptr) -> !llvm.ptr, i8
  call @mosaic_gpu_init_tma_desc (%33, ...)
  ...
  %44 = llvm.getelementptr %11[128] : (!llvm.ptr) -> !llvm.ptr, i8
  call @mosaic_gpu_init_tma_desc (%44, ...)
  ...
  %55 = gpu.launch async
    ^bb0:
      %66 = builtin.unrealized_conversion_cast_op(%22)
             : (!llvm.array<256 x i8>) -> !llvm.ptr


The type of a variable.

  Variables here are either operands or results of MLIR operations.


A unique identifier for a variable.


Hints are used to model propagation of layouts across operations.

  Since using `relayout`s is always an option in principle, propagation across
  ops can not rely only on an equation system. Instead, we introduce hints as
  an equation-like form of "soft constraints", i.e., it suggests that
  `variable` should be equal to `expression`.


Represents the way a shape is laid out in TMEM.

  The layout describes how the shape is split across the 128 rows (lanes) of
  TMEM. We reinterpret warp_dims as the partitioning of TMEM into 4 banks, each
  accessible from a single warp. The 32 lanes inside each bank are assigned
  consecutive elements from lane_dims. The data within each lane is linearized
  in row-major order, with each vector padded up to 32 bits (wider vectors are
  unsupported).


A tiling expression describing a permutation of elements of an nd-array.

  To apply one level of tiling to an array, each of the trailing dimensions (up
  to the rank of the tile) is unfolded into two dimensions: first equal to the
  ratio of the dimension size and the tile size, and second equal to the tile
  size. Then, all newly unfolded minor dimensions are transposed to appear at
  the end.

  This expression describes multi-level tiling, by applying each element of
  `tiles` in sequence to the array.

  See https://openxla.org/xla/tiled_layout for a more detailed explanation.


A FragmentedArray layout derived from a tiling expression.

  A logical array is transformed according to the tiling expression, and then
  split across warps (within a warpgroup), lanes, and vectorized according to
  the dimension indices. All dimension indices must be negative and should refer
  to the dimensions after tiling is applied.

  To better understand this layout, consider the example of WGMMA-related tiling
  from https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d as
  applied to a 128x128 array. The corresponding TiledLayout has a tiling of:

      (64, 8)(16, 8)(8, 8)(1, 2)

  and warp_dims=(-8,), lane_dims=(-4, -3), vector_dim=-1.

  We begin by applying the tiling (note that it always applies to a suffix):

          Tiled shape                       Remaining tiling actions
  ===========================================================================
  128 128                                  (64, 8)(16, 8)(8, 8)(1, 2)
    2  16  64  8                           (16, 8)(8, 8)(1, 2)
    2  16   4  1  16  8                    (8, 8)(1, 2)
    2  16   4  1   2  1  8  8              (1, 2)
    2  16   4  1   2  1  8  4  1  2

  The last expression is our final shape. At this stage, we're ready to partition
  the dimensions: warp_dims=(-8,) means that the 8-th dimension from the
  end is partitioned over 4 warps in a warpgroup (and so it must be of size 4).
  lane_dims=(-4, -3) indicate that those two dimensions are partitioned over
  the lanes within a warp (their product must be equal to 32, i.e. warp size).
  Finally, vector_dim=-1 indicates that each (logical) register is a vector
  containing 2 elements (there are no shape restrictions here).

  Given the above, the shape of the (logical) register array used to represent
  the array in each thread is: (2, 16, 1, 1, 2, 1, 1, 1, 1, 1). We have set all
  the dimensions above to 1, since each thread is a member of a single warp,
  a single lane, and the elements along the vectorized dimension are represented
  by a single (logical) register.


A fragmented array where all the values are equal represented as a register per thread.

  FragmentedArrays in this layout can be are always the result of a
  splat, each thread in the warpgroup has a single copy of the value,
  while the FragmentedArray pretends it has whatever shape the user
  wants. This means we can trivially broadcast, reshape and do
  elementwise operations with all other layouts.

  Examples:

  To load a value in
  ```
  FragmentedArray.splat(memref.load(ref_1d, [1]), (10,20,2))
  ```

  A shape is always provided for sanity check reasons.


Convert the array to 1D and then shard across threads.


Traversal orders with respect to the data flow for IR.


Semantics for the kernel's instruction stream.


A FragmentedArray that has is synchronized with the async proxy.

  This implies that it requires no additional synchronization when passed in
  as a WGMMA accumulator. In particular, when created from a
  FragmentedArray, the necessary synchronization is inserted at construction.


A variable is an abstract identifier.

  `key` is supposed to be hashable.


States that `source` must be relayout-able to `target`.

  Relayout-ability here is not defined as a fundamental property of layouts, but
  rather a reflection of our implementation. For instance, when evaluating this
  constraint, we will return `False` systematically if a relayout exists but we
  do not ever plan to support it.

  Modeling this constraint this way is helpful, in order to allow pruning
  inefficient solutions when attempting to solve an equation system.


States that `lhs != rhs`.


An equation system contains a set of equations and assignments.

  Assignments assign constant values to variables in the system (bound
  variables). Equations describe relationships between variables, and can be
  used to determine assignments for unknown (free) variables.

  Constraints are used to check predicates that must hold for the assignments to
  be valid.


Default WGMMA implementation.

  The kernel can accept any class that satisfies the same interface as this
  class.


Implementation details for jax2tf serialization.

  Abstract superclass for subclassing.


Tracer class that boxes a TF value and a JAX abstract value.

  In addition to the TF value we carry the JAX abstract value because
  there are some cases when it cannot be recovered from the value:
  when we are converting with polymorphic shapes or when the JAX aval
  has a custom element type. In these cases the shape of the value may
  have dimensions set to `None`, or it may only correspond to the JAX
  "physical" (TF/lowering-compatible) shape, so the JAX abstract value
  may contain more precise information.

  When the value has a partially-known shape, the dimensions marked as `None`
  must correspond to non-constant dimensions in the abstract value.

  See README.md for details.


Trace class that underlies the jax2tf transformation.

  We are going to ensure that jax2tf.convert is never nested inside other
  transformations. This is sufficient for intended use cases (converting
  fully-transformed JAX code). It also simplifies our job because we do not have
  to handle situations where we apply primitives on a mix of TF values and
  JAX tracers from an outer transformation. E.g., for addition both the TF
  values
  and the JAX tracers have an override and they get confused if they see values
  from the other world.

  Hence a TFT trace does not interact with non-TFT traces at lower-level. For
  higher-order control-flow primitives we invoke recursively
  _interpret_fun on the body of the conditional, which will create a nested TFT.

  We do want to allow transformations nested inside a TensorFlowTrace (TFT), but
  those will introduce their own MainTrace, and any operations involving those
  will be done on those traces, i.e., not a concern for TFT.


Tests that inspect the HLO for the sharding annotations.


Reloading output of jax2tf into JAX with call_tf


Reloading output of call_tf into TF with jax2tf.


Abstracts a few IO operation over standard "open" vs. gfile.


Tests a function with shape polymorphism.

  Converts `fun` with shape polymorphism, creates a `tf.ConcreteFunction`
  given `input_signature` and checks the inferred output shapes to match
  `expected_output_shapes`, then checks that the JAX and the TF functions
  produce the same results.


Tests for primitives that take shape values as parameters.


Compatibility tests that use TF.

  Uses tf.Graph to serialize and run the functions; expects that `func`
  contains a `jax2tf.call_tf` and uses `jax2tf.convert` to generate a
  `tf.Graph` containing a XlaCallModule with the actual MLIR module.


Specific primitive limitations for jax2tf.

  See the primitive_test module docstring for details.


A multi-layer perceptron.


A complete Graph Network model defined with Jraph.


A Graph Convolution Network + Pooling model defined with Jraph.


EncoderLSTM Module wrapped in a lifted scan transform.


LSTM encoder, returning state after finding the EOS token in the input.


DecoderLSTM Module wrapped in a lifted scan transform.

  Attributes:
    teacher_force: See docstring on Seq2seq module.
    vocab_size: Size of the vocabulary.


LSTM decoder.

  Attributes:
    init_state: [batch_size, hidden_size]
      Initial state of the decoder (i.e., the final state of the encoder).
    teacher_force: See docstring on Seq2seq module.
    vocab_size: Size of the vocabulary.


Sequence-to-sequence class using encoder/decoder architecture.

  Attributes:
    teacher_force: whether to use `decoder_inputs` as input to the decoder at
      every step. If False, only the first input (i.e., the "=" token) is used,
      followed by samples taken from the previous output logits.
    hidden_size: int, the number of hidden dimensions in the encoder and decoder
      LSTMs.
    vocab_size: the size of the vocabulary.
    eos_id: EOS id.


Bottleneck ResNet block.


Class defining the actor-critic model.


Applies word dropout to a batch of input IDs.

  This is basically the same as `nn.Dropout`, but allows specifying the
  value of dropped out items.


Embeds batches of token IDs into feature space.

  Attributes:
    vocab_size: The size of the vocabulary (i.e., the number of embeddings).
    embedding_size: The dimensionality of the embeddings.
    embedding_init: The initializer used to initialize the embeddings.
    frozen: Freezes the embeddings table, keeping it fixed at initial values.
    dropout_rate: Percentage of units to drop after embedding the inputs.
    word_dropout_rate: Percentage of input words to replace with unk_idx.
    unk_idx: The index (integer) to use to replace inputs for word dropout.


A simple unidirectional LSTM.


A simple bi-directional LSTM.


A simple Multilayer perceptron with 1 hidden layer.

  Attributes:
    hidden_size: The size of the hidden layer.
    output_size: The size of the output.
    activation: The activation function to apply to the hidden layer.
    dropout_rate: The dropout rate applied to the hidden layer.
    output_bias: If False, do not use a bias term in the last layer.
    deterministic: Disables dropout if set to True.


A classifier that uses attention to summarize the inputs.

  Attributes:
    hidden_size: The hidden size of the MLP classifier.
    output_size: The number of output classes for the classifier.
    dropout_rate: The dropout rate applied over the encoded_inputs, the summary
      of the inputs, and inside the MLP. Applied when `deterministic` is False.
    deterministic: Disables dropout if True.


A Text Classification model.


Adds (optionally learned) positional embeddings to the inputs.

  Args:
    config: TransformerConfig dataclass containing hyperparameters.
    decode: whether to run in single-position autoregressive mode.


Transformer MLP / feed-forward block.

  Args:
    config: TransformerConfig dataclass containing hyperparameters.
    out_dim: optionally specify out dimension.


Transformer encoder-decoder layer.

  Args:
    config: TransformerConfig dataclass containing hyperparameters.


Transformer Model Decoder for sequence to sequence translation.

  Args:
    config: TransformerConfig dataclass containing hyperparameters.
    shared_embedding: a shared embedding layer to use.


Transformer pure decoder stack for language modelling.

  Args:
    config: TransformerConfig dataclass containing hyperparameters.


Adds (optionally learned) positional embeddings to the inputs.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.


Transformer MLP / feed-forward block.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
    out_dim: optionally specify out dimension.


Transformer encoder layer.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.


Transformer Model for sequence tagging.


Adds (optionally learned) positional embeddings to the inputs.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
    decode: whether to run in single-position autoregressive mode.


Transformer encoder-decoder layer.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.


Transformer Model Encoder for sequence to sequence translation.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
    shared_embedding: a shared embedding layer to use.


Transformer Model Decoder for sequence to sequence translation.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
    shared_embedding: a shared embedding layer to use.


Transformer Model for sequence to sequence translation.

  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.


An MNIST model written using pure JAX.

  There is an option for the model to skip the classifier layer, for
  demonstrating reuse of the classifier-less model into a larger model.
  See README.md.


An MNIST model using Flax.


A simple CNN model for MNIST.

    There is an option for the model to skip the classifier layer, for
    demonstrating reuse of the classifier-less model into a larger model.
    See README.md.


A generalization of jax.Array that also supports fully remote arrays.


Experimental batched CSR matrix implemented in JAX.


Experimental COO matrix implemented in JAX.

  Note: this class has minimal compatibility with JAX transforms such as
  grad and autodiff, and offers very little functionality. In general you
  should prefer :class:`jax.experimental.sparse.BCOO`.

  Additionally, there are known failures in the case that `nse` is larger
  than the true number of nonzeros in the represented matrix. This situation
  is better handled in BCOO.


Environment for sparse jaxpr evaluation.

  The environment is essentially a collection of buffers and/or tracers
  that may be shared between one or more SparsifyValue objects, which
  represent sparse or dense arrays via indices into the list of buffers.


Experimental CSR matrix implemented in JAX.

  Note: this class has minimal compatibility with JAX transforms such as
  grad and autodiff, and offers very little functionality. In general you
  should prefer :class:`jax.experimental.sparse.BCOO`.

  Additionally, there are known failures in the case that `nse` is larger
  than the true number of nonzeros in the represented matrix. This situation
  is better handled in BCOO.


Experimental CSC matrix implemented in JAX; API subject to change.


Base class for high-level JAX sparse objects.


Experimental batched COO matrix implemented in JAX

  Args:
    (data, indices) : data and indices in batched COO format.
    shape : shape of sparse array.

  Attributes:
    data : ndarray of shape ``[*batch_dims, nse, *dense_dims]`` containing the
      explicitly stored data within the sparse matrix.
    indices : ndarray of shape ``[*batch_dims, nse, n_sparse]`` containing the
      indices of the explicitly stored data. Duplicate entries will be summed.

  Examples:
    Create a sparse array from a dense array:

    >>> M = jnp.array([[0., 2., 0.], [1., 0., 4.]])
    >>> M_sp = BCOO.fromdense(M)
    >>> M_sp
    BCOO(float32[2, 3], nse=3)

    Examine the internal representation:

    >>> M_sp.data
    Array([2., 1., 4.], dtype=float32)
    >>> M_sp.indices
    Array([[0, 1],
           [1, 0],
           [1, 2]], dtype=int32)

    Create a dense array from a sparse array:

    >>> M_sp.todense()
    Array([[0., 2., 0.],
           [1., 0., 4.]], dtype=float32)

    Create a sparse array from COO data & indices:

    >>> data = jnp.array([1., 3., 5.])
    >>> indices = jnp.array([[0, 0],
    ...                      [1, 1],
    ...                      [2, 2]])
    >>> mat = BCOO((data, indices), shape=(3, 3))
    >>> mat
    BCOO(float32[3, 3], nse=3)
    >>> mat.todense()
    Array([[1., 0., 0.],
           [0., 3., 0.],
           [0., 0., 5.]], dtype=float32)


A class representing detector and beam properties.


A class representing the half wave plate, as a transfermatrix stack on which operations are done


Parent class for MPI related stuff


Initialize a telescope and specify its properties.


    Given an instrument, create a scanning pattern and
    scan the sky.


    Represents the properties of a given material,
    possibly a unixial crystal.


    An object to hold all the properties of a stack, including thicknesses,
    angles of any uniaxial crystals' extraordinary axes, and material objects
    for all the layer materials, as defined above.


    Holds the transfer matrix for a single layer, along with
    the properties of the layer in a material object, defined above. This is
    calculated for a single frequency, index of refraction, and rotation of
    the layer about the z axis. This last rotation is only important for
    an anisotropic material.


    Calculates the transfer matrix for a stack, and creates
    an object which holds the individual layers plus the matrix for the stack as a whole.


    Class which represents the Geometry for a radiative transfer calculation.  Currently this is made up of
    a set of :py:class:`sasktran.LineOfSight` which define the viewing geometry, and a :py:attr:`sasktran.Geometry.sun`
    property which controls the solar position.  In the future this class may be extended to also control options
    such as the top of the atmosphere, and the 'ground' altitude.

    Note that it is not always necessary to explicitly set the sun position.  If no sun position is given then it
    is automatically calculated from the average `mjd` given in the lines of sight.

    This class can be thought of as the geometry interface to the radiative transfer model.  For simple
    problems, we can construct this class and modify the lines of sight directly, as is done in the example below.
    For more complicated problems it may be desirable to create a new class that inherits from this class, as is
    done in the :class:`VerticalImage` class.

    Examples
    --------
    >>> import sasktran as sk
    >>> geometry = sk.Geometry()
    >>> los1 = sk.LineOfSight(mjd=54832.5, observer=[3.6760131547888e+005, 1.0099763136400e+006, -6.871601202127e+006],\
                              look_vector=[2.884568631765662e-001, 7.925287180643269e-001,  5.372996083468238e-001]),
    >>> los2 = sk.LineOfSight(mjd=54832.5, observer=[3.6928085406796e+005, 1.0145908079886e+006, -6.870844156040e+006],\
                              look_vector=[2.88456863176566e-001, 7.925287180643269e-001,  5.372996083468238e-001])
    >>> geometry.lines_of_sight = [los1, los2]
    >>> print(geometry)
    Geometry object containing 2 lines of sight
    >>> geometry.sun = [0, 0, 1]
    >>> print(geometry)
    Geometry object containing 2 lines of sight
    Sun position forced to [0, 0, 1]


    A specialized :class:`Geometry` class which has various convenience methods to construct lines of sight
    corresponding to a vertical image of the atmosphere based on input solar zenith angles.

    Examples
    --------
    >>> from sasktran.geometry import VerticalImage
    >>> geometry = VerticalImage()
    >>> geometry.from_sza_saa(60, 60, 0, 0, [10, 20, 30, 40], 54372, 0)
    >>> print(geometry)
    Geometry object containing 4 lines of sight
    Sun position forced to [ 0.5        0.75       0.4330127]


    Class which implements a Bidirectional Reflectance Distribution Function.  This is a light wrapper around the
    sasktranif ISKBrdf object.  It is recommended to use the specialized BRDF classes rather than this one.

    A BRDF is defined as the ratio of the radiance scattered from a surface into a given direction to
    the collimated power incident on a unit area of the surface. Much of the literature on BRDFs (Spurr 2002, Roujean
    1992, Rahman 1993, Wanner 1995, Cox and Munk 1954, Kokhanovsky 2012) define reflectance functions through a
    quantity that is scaled up by pi compared to this definition.

    Sasktran defines the relative azimuth angle such that 0 degrees corresponds to forward scattering. Much of the
    literature on BRDFs (Spurr 2002, Roujean 1992, Rahman 1993, Wanner 1995, Cox and Munk 1954, Hapke 2012) define it
    the opposite way.

    Parameters
    ----------
    name: str
        Name of the BRDF object to create

    Examples
    --------
    >>> from sasktran import BRDF
    >>> brdf = BRDF('lambertian')
    >>> brdf.skif_object().SetProperty('Albedo', 0.3)
    True
    >>> brdf.reflectance(340, 0, 0, 54372, 0.5, 0.5, 0)
    0.0954929658551372


    Class which implements a Lambertian Bidirectional Reflectance Distribution Function.

    Parameters
    ----------
    albedo: float
        Ratio of incoming to outgoing radiance.  The reflectance distribution is equal to albedo/pi

    Examples
    --------
    >>> from sasktran.brdf import Lambertian
    >>> brdf = Lambertian(albedo=0.3)
    >>> brdf.reflectance(340, 0, 0, 54372, 0.5, 0.5, 0)
    0.0954929658551372
    >>> brdf.albedo = 1
    >>> brdf.reflectance(340, 0, 0, 54372, 0.5, 0.5, 0)
    0.3183098861837907


    Implementation of the snow BRDF from Kokhanovsky 2012.  The parameters L and M are optional, the default values
    correspond to 1020 nm over greenland.

    Equation 1 is implemented with the following changes:

    * the reflectance is divided by pi

    Parameters
    ----------
    L: float, optional
        Parameter approximately equal to 13 times the average snow grain optical diameter. Default 3.6e6

    M: float, optional
        Parameter proportional to the mass concentration of pollutants in the snow. Default 5.5e-8

    Examples
    --------
    >>> from sasktran.brdf import Kokhanovsky
    >>> brdf = Kokhanovsky()
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.25185744490707646

    References
    ----------
    A. A. Kokhanovsky and F. M. Breon, "Validation of an Analytical Snow BRDF Model
    Using PARASOL Multi-Angular and Multispectral Observations,"
    IEEE Geoscience and Remote Sensing Letters, vol. 9, no. 5, pp. 928-932, Sept. 2012.
    doi: 10.1109/LGRS.2012.2185775
    URL: http://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=6166850&isnumber=6205680


    Implementation of a vegetation BRDF based on Hapke 2012 and Spurr 2002.

    Equation A.10 of Spurr is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Parameters
    ----------
    omega: float
        Single scatter albedo
    delta: float
        Angular width at half maximum of the hotspot / opposition peak
    b0: float
        Amplitude of the hotspot / opposition peak

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.Hapke(omega=0.6, delta=0.06, b0=1.0)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.0814968276380363

    References
    ----------
    Hapke B. Theory of reflectance and emittance spectroscopy. Cambridge:
    Cambridge University Press, 2012.

    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832


    Implementation of the Ross-thick Li-sparse-reciprocal BRDF used in MODIS' Albedo/BRDF product. It is a linear
    combination of an isotropic (Lambertian) term, a volume scatterer (Ross-thick kernel) and a geometric scatterer
    (Li-sparse-reciprocal kernel).

    Parameters
    ----------
    f_iso: float
        Weight of the isotropic (Lambertian) kernel
    f_vol: float
        Weight of the volume (Ross-thick) kernel
    f_geo: float
        Weight of the geometric (Li-sparse-reciprocal) kernel

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.MODIS(f_iso=0.2, f_vol=0.2, f_geo=0.1)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.031577469878097834

    References
    ---------
    A. H. Strahler, J. P. Muller, "MODIS BRDF/Albedo Product: Algorithm
    Theoretical Basis Document Version 5.0," April 1999.
    URL: https://lpdaac.usgs.gov/sites/default/files/public/product_documentation/atbd_mod09_v5.pdf


    Implements a BRDF that is a linear combination of an arbitrary number of kernels.

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.LinearCombination()

    >>> kernel1 = sk.Lambertian(albedo=1.0)
    >>> kernel2 = sk.RossThickKernel()
    >>> kernel3 = sk.LiSparseReciprocalKernel(crown_shape=1.0, relative_height=2.0)

    >>> brdf.kernels = [kernel1, kernel2, kernel3]
    >>> brdf.kernel_weights = [0.2, 0.2, 0.1]

    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.031577469878097834


    Base class for Li kernels from Wanner 1995 and Spurr 2002. It is recommended to use derived classes
    (LiSparseKernel, LiDenseKernel or LiSparseReciprocalKernel) instead of this class. It is also recommended to use
    kernels as part of a LinearCombination and not as standalone BRDFs.

    Tree crowns are modelled as spheroids of height 2b and width 2r, with centers a height of h above the ground.

    Parameters
    ----------
    crown_shape: float
        The ratio b/r describing the shape of the crown.
    relative_height: float
        The ratio h/b describing the height of the crown above the ground.

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.LiKernel('li_sparse_kernel', crown_shape=1.0, relative_height=2.0)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    -0.8753521870054244

    References
    ----------
    W. Wanner and A. H. Strahler, "On the derivation of kernels for kernel-driven
    modes of bidirectional reflectance," Journal of Geophysical Research, vol.
    100, no. d10, pp. 21077-21089, Oct. 20, 1995.
    doi: 10.1029/95JD02371
    URL: http://onlinelibrary.wiley.com/doi/10.1029/95JD02371/abstract

    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832


    The Li-sparse kernel models sparse canopy where mutual shadowing can be ignored. It is recommended to use
    kernels as part of a LinearCombination and not as standalone BRDFs.

    Equation 32 from Wanner is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Parameters
    ----------
    crown_shape: float, optional
        See sk.LiSparse - Default 1.0
    relative_height: float, optional
        See sk.LiSparse - Default 2.0

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.LiSparseKernel(crown_shape=1.0, relative_height=2.0)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    -0.8753521870054244

    References
    ----------
    W. Wanner and A. H. Strahler, "On the derivation of kernels for kernel-driven
    modes of bidirectional reflectance," Journal of Geophysical Research, vol.
    100, no. d10, pp. 21077-21089, Oct. 20, 1995.
    doi: 10.1029/95JD02371
    URL: http://onlinelibrary.wiley.com/doi/10.1029/95JD02371/abstract

    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832


    The Li-sparse-reciprocal kernel is a variation on the Li-sparse kernel which allows the incident/viewing angles
    to be swapped without changing the BRDF's value (reciprocal) - See the MODIS Albedo/BRDF ATBD. It is recommended
    to use kernels as part of a LinearCombination and not as standalone BRDFs.

    Equation 39 from Strahler is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Parameters
    ----------
    crown_shape: float, optional
        See sk.LiSparse - Default 1.0
    relative_height: float, optional
        See sk.LiSparse - Default 2.0

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.LiSparseReciprocalKernel(crown_shape=1.0, relative_height=2.0)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    -0.477464829275686

    References
    ----------
    A. H. Strahler, J. P. Muller, "MODIS BRDF/Albedo Product: Algorithm
    Theoretical Basis Document Version 5.0," April 1999.
    URL: https://lpdaac.usgs.gov/sites/default/files/public/product_documentation/atbd_mod09_v5.pdf


    The Li-dense kernel models dense canopy where mutual shadowing cannot be ignored. It is recommended to use
    kernels as part of a LinearCombination and not as standalone BRDFs.

    Equation 47 from Wanner is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Parameters
    ----------
    crown_shape: float, optional
        See sk.LiSparse - Default 2.5
    relative_height: float, optional
        See sk.LiSparse - Default 2.0

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.LiDenseKernel(crown_shape=2.5, relative_height=2.0)
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    -0.46940635114445084

    References
    ----------
    W. Wanner and A. H. Strahler, "On the derivation of kernels for kernel-driven
    modes of bidirectional reflectance," Journal of Geophysical Research, vol.
    100, no. d10, pp. 21077-21089, Oct. 20, 1995.
    doi: 10.1029/95JD02371
    URL: http://onlinelibrary.wiley.com/doi/10.1029/95JD02371/abstract

    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832


    The Ross-thin kernel models canopies with a small leaf area index (LAI) - see Wanner 1995 and Spurr 2002. It is
    recommended to use kernels as part of a LinearCombination and not as standalone BRDFs.

    Equation 13 from Wanner is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.RossThinKernel()
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.8132395113781661

    References
    ----------
    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832

    W. Wanner and A. H. Strahler, "On the derivation of kernels for kernel-driven
    modles of bidirectional reflectance," Journal of Geophysical Research, vol.
    100, no. d10, pp. 21077-21089, Oct. 20, 1995.
    doi: 10.1029/95JD02371
    URL: http://onlinelibrary.wiley.com/doi/10.1029/95JD02371/abstract


    The Ross-thick kernel models canopies with a large leaf area index (LAI) - see Wanner 1995 and Spurr 2002. It is
    recommended to use kernels as part of a LinearCombination and not as standalone BRDFs.

    Equation 7 from Wanner is implemented with the following changes:

    * cos(phi) is replaced with -cos(phi)
    * the reflectance is divided by pi

    Examples
    --------
    >>> import sasktran as sk
    >>> brdf = sk.RossThickKernel()
    >>> brdf.reflectance(1050, 0, 0, 54372, 0.5, 0.5, 0)
    0.07830987784454152

    References
    ----------
    Robert J. D. Spurr, "A new approach to the retrieval of surface properties
    from earthshine measurements," Journal of Quantitative Spectroscopy
    & Radiative Transfer, vol. 83, pp. 15-46, Oct. 9, 2002.
    doi:10.1016/S0022-4073(02)00283-2
    url: http://www.sciencedirect.com/science/article/pii/S0022407302002832

    W. Wanner and A. H. Strahler, "On the derivation of kernels for kernel-driven
    modles of bidirectional reflectance," Journal of Geophysical Research, vol.
    100, no. d10, pp. 21077-21089, Oct. 20, 1995.
    doi: 10.1029/95JD02371
    URL: http://onlinelibrary.wiley.com/doi/10.1029/95JD02371/abstract


    A custom BRDF that is constructed as the combination of other BRDF's on a uniform latitude/longitude grid.

    Parameters
    ----------
    longitudes: np.array
        Longitude grid that the BRDF is specified on in degrees.  Longitudes must be specified from 0 to 360
        and not from -180 to 180.
    latitudes: np.array
        Latitude grid that the BRDF is specified on in degrees
    brdfs: np.ndarray
        2-D Array of dimension (len(longitudes), len(latitudes)) of sk.BRDF objects

    Examples
    --------
    >>> import sasktran as sk
    >>> latitudes = [-10, 10]
    >>> longitudes = [40, 60]
    >>> brdfs = np.empty((len(longitudes), len(latitudes)), dtype=object)
    >>> brdfs[0, 0] = sk.Lambertian(0.1)
    >>> brdfs[0, 1] = sk.Lambertian(0.5)
    >>> brdfs[1, 0] = sk.Lambertian(0.8)
    >>> brdfs[1, 1] = sk.Lambertian(1.0)
    >>> brdf = sk.LatLonBRDF(longitudes, latitudes, brdfs)
    >>> brdf.reflectance(wavel_nm=750, latitude=0, longitude=50, mjd=54372, mu_in=0, mu_out=0, cos_dphi=0)
    0.1909859317102744


    A custom BRDF that is composed of a collection of other BRDFs on a wavelength grid to create a spectrally
    varying BRDF.

    Parameters
    ----------
    wavelengths: np.array
        Wavelength grid that the BRDF is specified on in nm
    brdfs: np.ndarray
        1-D Array of dimension (len(wavelengths)) of sk.BRDF objects

    Examples
    --------
    >>> import sasktran as sk
    >>> import numpy as np
    >>> wavelengths = np.array([350, 750])
    >>> brdfs = np.empty((len(wavelengths)), dtype=object)
    >>> brdfs[0] = sk.Lambertian(0.1)
    >>> brdfs[1] = sk.Lambertian(0.5)
    >>> brdf = sk.SpectrallyVaryingBRDF(wavelengths, brdfs)
    >>> brdf.reflectance(wavel_nm=350, latitude=0, longitude=50, mjd=54372, mu_in=0, mu_out=0, cos_dphi=0)
    0.03183098861837907
    >>> brdf.reflectance(wavel_nm=550, latitude=0, longitude=50, mjd=54372, mu_in=0, mu_out=0, cos_dphi=0)
    0.09549296585513721


    Defines a stokes vector with its associated basis.

    Parameters
    ----------
    stokes : numpy array shape (4,)
        The stokes vector [I, Q, U, V]

    basis : numpy array shape (3,3)
        Coordinate basis the stokes vector is defined in.  basis[0, :] is the propagation direction, basis[1, :] is
        the theta direction, and basis[2, :] is the phi direction.  Directions are specified in ECEF coordinates.
        The basis must be constructed such that basis[1, :] cross basis[2, :] is equal to basis[0, :]

    Raises
    ------
    ValueError
        If the basis is not correctly constructed

    Examples
    --------
    >>> from sasktran import StokesVector
    >>> basis = np.eye(3, 3)
    >>> stokes_rad = [1, 0.1, -0.1, 0]
    >>> stokes_v = StokesVector(stokes_rad, basis)
    >>> print(stokes_v)
    Stokes Vector
    I: 1 Q: 0.1 U: -0.1 V: 0
    Propagation: [ 1.  0.  0.]
    Theta: [ 0.  1.  0.]
    Phi: [ 0.  0.  1.]


    The Atmosphere class is a generic container that defines the atmospheric state for a radiative transfer calculation.
    The first component of this is defining a set of atmospheric constituents, which is controlled by the
    :py:attr:`sasktran.Atmosphere.species` property.

    The second component is defining the ground reflectance. Sasktran support :ref:`BRDF` surfaces. The
    :ref:`BRDF` can be set with the :py:attr:`sasktran.Atmosphere.brdf` property.

    Emissions can also be added to the atmosphere using the :py:attr:`sasktran.Atmosphere.brdf` property.  Note that
    currently the only radiative transfer model that supports emissions is the HR engine.

    Additionally, the atmosphere object can be used to set the temperature and pressure :ref:`climatology` (see
    :py:attr:`sasktran.Atmosphere.atmospheric_state`), as well as species to calculate weighting functions for (see
    :py:attr:`sasktran.Atmosphere.wf_species`).


    A climatology quantifying something in/about the atmosphere. For example climatologies could be used
    to define profiles for temperature, pressure, gas number densities, etc.

    Parameters
    ----------
    name : str
        Name of the climatology to create

    Examples
    --------
    >>> import sasktran as sk
    >>> msis = sk.Climatology('msis90')
    >>> print(msis)
    SasktranIF Climatology: msis90
    Supported Species: ['SKCLIMATOLOGY_PRESSURE_PA', 'SKCLIMATOLOGY_AIRNUMBERDENSITY_CM3', 'SKCLIMATOLOGY_TEMPERATURE_K', 'SKCLIMATOLOGY_O2_CM3']


    A special climatology which handles user defined values.

    Parameters
    ----------
    altitudes: np.ndarray
        Array of altitudes in meters that the profile is to be specified at.  Same shape as values.
    values: dict
        dictionary of species, value pairs. Values should be the same length as altitudes.
    interp: str, optional
        One of 'linear' or 'log'. Defines the interpolation space. Default is 'linear'
    spline: bool, optional
        One of True of False.  If True then a quadratic spline will be used to interpolate, if False then piecewise
        linear interpolation is done.  Default is False.

    Examples
    --------
    >>> from sasktran.climatology import ClimatologyUserDefined
    >>> import numpy as np
    >>> altitudes = np.linspace(500, 99500, 100)
    >>> values = np.ones_like(altitudes)
    >>> climatology = ClimatologyUserDefined(altitudes, {'SKCLIMATOLOGY_O3_CM3': values})
    >>> climatology.get_parameter('SKCLIMATOLOGY_O3_CM3', latitude=0, longitude=0, altitudes=[10000, 11000], mjd=54372)
    array([ 1.,  1.])
    >>> climatology['SKCLIMATOLOGY_O3_CM3'] *= 2
    >>> climatology.get_parameter('SKCLIMATOLOGY_O3_CM3', latitude=0, longitude=0, altitudes=[10000, 11000], mjd=54372)
    array([ 2.,  2.])


    A two dimensional user defined climatology in altitude and angle within a user specified plane.

    Parameters
    ----------
    angles_deg: np.ndarray
        One dimensional array of the angular grid in degrees
    alts_m: np.ndarray
        One dimensional array of the height grid in meters
    values: dict
        dictionary of species, value pairs. Values should be two dimensional arrays of shape (len(angles_deg), len(alts_m))
    reference_vector: np.ndarray
        Length 3 unit vector defining the location where the angle is 0, i.e., the x-axis of the plane
    normal_vector: np.ndarray
        Length 3 unit vector defining the normal vector of the plane.


 A user defined climatology class for profiles which have a geographic dependence.

    For more information see `sk.ClimatologyUserDefined`.

    Parameters
    ----------
    lats : np.ndarray
        The latitude grid (1D) for the climatology. The latitudes for the region will be queried must be defined.
        Latitudes don't need to be evenly spaced.
    lons : np.ndarray
        The longitude grid (1D) for the climatology. Must span 0 to 360 longitude (but doesn't need even spacing).
    alts : np.ndarray
        The altitude grid for the climatology in meters. Altitudes don't need to be evenly spaced. Altitudes at the
        top of that atmosphere (TOA) and the ground (elevation) must be defined.
    values:
        A 3D array of values with dimensions: latitude, longitude, altitude.
    interp: str, optional
        One of 'linear' or 'log'. Defines the interpolation space. Default is 'linear'
    spline: bool, optional
        One of True of False.  If True then a quadratic spline will be used to interpolate, if False then piecewise
        linear interpolation is done.  Default is False.


    The Labow volume mixing ratio climatology of ozone. This is a climatology of the volume mixing ratio from
    0 km to 60 km in 1 km increments for 18 latitudes in steps of 10 degrees from -85 to +85 for each month of the year.
    The climatology seems to be unpublished as all references refer to an unpublished piece of work by McPeters and
    Labow in 2002 or 2003. We have extended the climatology above 60 km using the differences in scale height between
    the neutrals (7 km scale height ) and ozone (4.5 km scale height) to calculate a new scale height to extrapolate
    the signal at 60 km. We have also extended the volume mixing ratio so it can be converted to an ozone density.
    This requires an atmospheric number density climatology for which the MSIS90 climatology is used:


    A climatology for the NO2 number density based upon the Prather photo-chemical box model. This climatological
    model is based upon monthly mean solutions to the box model for a range of latitudes, local solar times and
    altitudes.


    A climatology of number density, pressure, and temperature based on the ERA interim reanalysis.  This climatology
    requires additional configuration, see :ref:`configuration` for more information.


    The global stratospheric aerosol extinction climatology at 8 wavelengths from 386-3400nm that spans from 1979-2016
    and incorporates data from numerous satellite datasets. Interpolations in latitude and time use nearest neigbour.

    Note about wavelengths:
        386, 452, 525 and 1020 are the merged aerosol extinction datasets, although 386 and 452nm are not generally
        recommended. 525 and 1020nm channels provide the most robust extinction values so for radiative transfer
        calculations it is recommended to use one of these two channels.

        Other channels (750, 1257 and 3400) are from individual instruments.

    As the GloSSAC climatology is extinction based the species_name is 'SKCLIMATOLOGY_AEROSOL_EXTINCTIONPERKM', although
    the skif_object is converted to number density for use in sasktran.

    Parameters
    ----------
    wavelength_nm: float
        wavelength of the returned extinction. Possible values are: 386, 452, 525, 750, 1020, 1257, 3400
    extend: bool
        If true the extinction profile is extended above and below the valid range to avoid sharp distcontinuities.
        Profiles are assumed to exponenetially decay outside the valid range. Default is False where all non-valid
        regions are set to zero. The class variables decay_below and decay_above can be used to adjust the decay rates
        below and above the valid range respectively


    Class which represents a single line of sight in SASKTRAN.  A single line of sight is defined by the observer
    position, a timestamp, and a unit look vector.

    Parameters
    ----------
    mjd : float
        Modified julian date for the measurement

    observer : np.array
        Three element array indicating the observer position for the measurement.

    look_vector : np.array
        Three element array which is the unit look vector away from the instrument

    Examples
    --------
    >>> from sasktran import LineOfSight
    >>> los = LineOfSight(mjd=54832.5, observer=[3.6760131547888e+005, 1.0099763136400e+006, -6.871601202127e+006],\
                          look_vector=[2.884568631765662e-001, 7.925287180643269e-001,  5.372996083468238e-001])
    >>> print(los)
    Observer: [367601.31547888, 1009976.31364, -6871601.202127]
    Look: [0.2884568631765662, 0.7925287180643269, 0.5372996083468238]
    MJD: 54832.5
    >>> print(los.mjd)
    54832.5
    >>> print(los.observer)
    [367601.31547888, 1009976.31364, -6871601.202127]
    >>> print(los.look_vector)
    [0.2884568631765662, 0.7925287180643269, 0.5372996083468238]


Defines the interface required to create an emission that can be used with a radiative transfer model.


    Class which represents the emission of a species in the atmosphere.  Most of the time it is not recommended to use
    this class and more specific classes, such as EmissionTable, should be used instead.

    Parameters
    ----------
    name : str
        Name handle of the emission to create.


    Defines an emission in the atmosphere represented with a table of volume emission rates on an altitude/wavelength
    grid.

    Parameters
    ----------
    altitude_m : np.ndarray
        Array of altitudes in [m] defining the altitude dimension of the emission grid

    wavel_nm : np.ndarray
        Array of wavelengths in [nm] defining the wavelength dimension of the altitude grid

    volumeemission : np.ndarray
        Two dimensional array of volume emission rates in [photons/nm/sec/ster] with dimension (len(alts), len(wavel))

    Examples
    --------
    >>> import sasktran as sk
    >>> alts = np.arange(500, 100500, 1000)
    >>> wavel = np.linspace(300, 350, 50)
    >>> ver = np.ones((len(alts), len(wavel)))
    >>> emission = sk.EmissionTable(alts, wavel, ver)
    >>> print(emission)
    Emission of type: USERDEFINED_WAVELENGTHHEIGHT


    Class used to enable thermal emissions within the HR radiative transfer model.  Currently only works as part
    of a full radiative transfer calculation.

    Examples
    --------
    >>> import sasktran as sk
    >>> emission = sk.EmissionThermal()
    >>> geometry = sk.VerticalImage()
    >>> geometry.from_sza_saa(sza=60, saa=60, lat=0, lon=0, tanalts_km=[10, 20, 30, 40], mjd=54372, locallook=0, \
        satalt_km=600, refalt_km=20)
    >>> atmosphere = sk.Atmosphere()
    >>> atmosphere['ozone'] = sk.Species(sk.O3OSIRISRes(), sk.Labow())
    >>> atmosphere['air'] = sk.Species(sk.Rayleigh(), sk.MSIS90())
    >>> atmosphere.emissions['thermal'] = emission
    >>> engine = sk.EngineHR(geometry, atmosphere, wavelengths=[1000])
    >>> engine.num_orders_of_scatter = 1
    >>> engine.include_emissions = True
    >>> engine.calculate_radiance()
    array([[ 0.01218222,  0.00270503,  0.00055809,  0.00013423]])


    An object representing some atmospheric constituent. A species is an object that associates a constituents
    :ref:`optical_property` and :ref:`climatology`.

    Examples
    --------
    >>> import sasktran as sk
    >>> no2 = sk.Species(optical_property=sk.NO2Vandaele1998(), climatology=sk.Pratmo())
    >>> print(no2)
    SKCLIMATOLOGY_NO2_CM3
    SasktranIF Climatology: NO2PRATMO
    Supported Species: ['SKCLIMATOLOGY_NO2_CM3']
    SasktranIF Optical Property: NO2_VANDAELE1998


    Special climatology which handles aerosol distributions. Contains a particle size Climatology and
    MieAerosol OpticalProperty internally to handle changes to size parameters.

    Parameters
    ----------
    altitudes: np.ndarray
        Array of altitudes in meters that the profile is to be specified at.  Same shape as values.
    aerosol_values: dict
        Dictionary of species names and values containing the aerosol number density Values should be the same size
        as altitudes.
    particle_size_values: dict
        Dictionary of species names and values containing the particle size profile. Values should be the same size
        as altitudes.
    species: str (optional, default='H2SO4')
        Molecule to use, one of ['H2SO4', 'ICE', 'WATER']
    interp: str, optional
        One of 'linear' or 'log'. Defines the interpolation space. Default is 'linear'
    spline: bool, optional
        One of True of False.  If True then a quadratic spline will be used to interpolate, if False then piecewise
        linear interpolation is done.  Default is False.

    Examples
    --------
    >>> from sasktran.species import SpeciesAerosol
    >>> import numpy as np
    >>> altitudes = np.linspace(500, 99500, 100)
    >>> values = np.ones_like(altitudes)
    >>> species = SpeciesAerosol(altitudes, {'SKCLIMATOLOGY_AEROSOL_CM3': values},\
                                            {'SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS': values * 0.08,\
                                             'SKCLIMATOLOGY_LOGNORMAL_MODEWIDTH': values * 1.6}, 'H2SO4')
    >>> species.get_parameter('SKCLIMATOLOGY_AEROSOL_CM3', latitude=0, longitude=0, altitudes=[10000, 11000], mjd=54372)
    array([ 1.,  1.])
    >>> species['SKCLIMATOLOGY_AEROSOL_CM3'] *= 2
    >>> species.get_parameter('SKCLIMATOLOGY_AEROSOL_CM3', latitude=0, longitude=0, altitudes=[10000, 11000], mjd=54372)
    array([ 2.,  2.])
    >>> species.get_parameter('SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS', latitude=0, longitude=0,\
                              altitudes=[10000, 11000], mjd=54372)
    array([ 0.08,  0.08])
    >>> species['SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS'] *= 2
    >>> species.get_parameter('SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS', latitude=0, longitude=0,\
                              altitudes=[10000, 11000], mjd=54372)
    array([ 0.16,  0.16])


    __init__(self, altitudes: np.ndarray=np.arange(0, 100001, 500), particle_size_values: dict=None, \
    species: str='H2SO4', interp: str='linear', spline: bool=False, extend=False)

    Special climatology which handles aerosol distributions. Similar to :py:class:`sasktran.SpeciesAerosol`, however
    rather than a user defined aerosol profile the :py:class:`sasktran.GloSSAC` climatology is used.

    A dictionary of particle size values can be provided. If not, a lognormal distribution with median radius=80nm and
    width of 1.6 is used.

    Notes
    -----
    Although the user side uses the handle 'SKCLIMATOLOGY_AEROSOL_EXTINCTIONPERKM', aerosol properties are
    passed to sasktran using the 'SKCLIMATOLOGY_AEROSOL_CM3' handle

    Examples
    --------
    >>> from sasktran.species import SpeciesAerosolGloSSAC
    >>> import numpy as np
    >>> altitudes = np.linspace(500, 99500, 100)
    >>> values = np.ones_like(altitudes)
    >>> species = SpeciesAerosolGloSSAC()
    >>> species.get_parameter('SKCLIMATOLOGY_AEROSOL_EXTINCTIONPERKM', latitude=0, longitude=0, \
                              altitudes=[22000, 23000], mjd=54372)
    array([ 0.00049428,  0.0004945 ])
    >>> species.get_parameter('SKCLIMATOLOGY_AEROSOL_EXTINCTIONPERKM', latitude=20, longitude=0, \
                              altitudes=[22000, 23000], mjd=54372)
    array([ 0.00048611,  0.0004257 ])
    >>> species.get_parameter('SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS', latitude=0, longitude=0, \
                               altitudes=[10000, 11000], mjd=54372)
    array([ 0.08,  0.08])
    >>> species['SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS'] *= 2
    >>> species.get_parameter('SKCLIMATOLOGY_LOGNORMAL_MODERADIUS_MICRONS', latitude=0, longitude=0, \
                               altitudes=[10000, 11000], mjd=54372)
    array([ 0.16,  0.16])