coordax 0.2.1__tar.gz → 0.2.2__tar.gz

{coordax-0.2.1 → coordax-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.1
+Version: 0.2.2
 Summary: Coordinate axes for scientific computing in JAX
 Author-email: Google LLC <noreply@google.com>
 License-Expression: Apache-2.0

{coordax-0.2.1 → coordax-0.2.2}/coordax/__init__.py

@@ -53,4 +53,4 @@ from coordax.ndarrays import (
 )
 import coordax.testing  # pylint: disable=unused-import
 
-__version__ = '0.2.1'  # keep sync with pyproject.toml
+__version__ = '0.2.2'  # keep sync with pyproject.toml

{coordax-0.2.1 → coordax-0.2.2}/coordax/coordinate_systems.py

@@ -46,7 +46,7 @@ Sequence = collections.abc.Sequence
 @functools.partial(utils.export, module='coordax.coords')
 @dataclasses.dataclass(frozen=True)
 class NoCoordinateMatch:
-  """For use when a coordinate does not match an xarray.Coordinates object."""
+  """For use when no Coordax coordinate matches xarray coordinate."""
 
   reason: str
 
@@ -56,13 +56,13 @@ class Coordinate(abc.ABC):
   """Abstract class for coordinate objects.
 
   Coordinate subclasses are expected to obey several invariants:
-  1. Dimension names may not be repeated: `len(set(dims)) == len(dims)`
-  2. All dimensions must be named: `len(shape) == len(dims)`
+  1. Dimension names may not be repeated: ``len(set(dims)) == len(dims)``
+  2. All dimensions must be named: ``len(shape) == len(dims)``
 
   Every non-abstract Coordinate subclass must be registered as a "static"
   pytree node, e.g., by decorating the class with
-  `@jax.tree_util.register_static`. Static pytrees nodes must implement
-  `__hash__` and `__eq__` according to the requirements of keys in Python
+  ``@jax.tree_util.register_static``. Static pytrees nodes must implement
+  ``__hash__`` and ``__eq__`` according to the requirements of keys in Python
   dictionaries. This is easiest to acheive with frozen dataclasses, but care
   must be taken when working with np.ndarray attributes.
   """
@@ -73,7 +73,7 @@ class Coordinate(abc.ABC):
     """Dimension names of the coordinate.
 
     All subclasses must return a tuple of dimension names as strings, with the
-    exception of `DummyAxis`.
+    exception of ``DummyAxis``.
     """
     raise NotImplementedError()
 
@@ -135,8 +135,12 @@ class Coordinate(abc.ABC):
         values.
 
     Returns:
-      A matching instance of this coordinate or `NoCoordinateMatch` if this
-      coordinate does not match the xarray dimensions and coordinates.
+      A matching instance of this coordinate or ``NoCoordinateMatch`` if this
+      coordinate type does not match the xarray dimensions and coordinates.
+
+    See also:
+      :func:`coordax.from_xarray`
+      :func:`coordax.coords.from_xarray`
     """
     raise NotImplementedError('from_xarray not implemented')
 
@@ -513,6 +517,7 @@ class LabeledAxis(Coordinate):
 @functools.partial(utils.export, module='coordax.coords')
 def compose(*coordinates: Coordinate) -> Coordinate:
   # pylint: disable=line-too-long
+  # fmt: off
   """Compose coordinates into a unified coordinate system.
 
   Args:
@@ -528,6 +533,7 @@ def compose(*coordinates: Coordinate) -> Coordinate:
     >>> cx.coords.compose(x, y)
     CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.SizedAxis('y', size=3)))
   """
+  # fmt: on
   product = CartesianProduct(coordinates)
   match len(product.coordinates):
     case 0:
@@ -544,7 +550,8 @@ def insert_axes(
     indices_to_axes: dict[int, Coordinate],
 ) -> Coordinate:
   # pylint: disable=line-too-long
-  """Returns `coordinate` with extra axes inserted at specified positions.
+  # fmt: off
+  """Returns ``coordinate`` with extra axes inserted at specified positions.
 
   Args:
     coordinate: The coordinate system to modify.
@@ -561,6 +568,7 @@ def insert_axes(
     >>> cx.coords.insert_axes(x, {1: z})
     CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.SizedAxis('z', size=4)))
   """
+  # fmt: on
   indices_to_axes = indices_to_axes.copy()
   ndim = coordinate.ndim + len(indices_to_axes)
   normalize_idx = lambda i: i + ndim if i < 0 else i
@@ -581,7 +589,8 @@ def replace_axes(
     replace_with: Coordinate,
 ) -> Coordinate:
   # pylint: disable=line-too-long
-  """Returns `coordinate` with `to_replace` replaced by `replace_with`.
+  # fmt: off
+  """Returns ``coordinate`` with ``to_replace`` replaced by ``replace_with``.
 
   Args:
     coordinate: The coordinate system to modify.
@@ -589,10 +598,10 @@ def replace_axes(
     replace_with: The new coordinate to insert.
 
   Returns:
-    A new coordinate object with `to_replace` substituted by `replace_with`.
+    A new coordinate object with ``to_replace`` substituted by ``replace_with``.
 
   Raises:
-    ValueError: If `to_replace` is not a contiguous part of `coordinate`.
+    ValueError: If ``to_replace`` is not a contiguous part of ``coordinate``.
 
   Examples:
     >>> import coordax as cx
@@ -605,6 +614,7 @@ def replace_axes(
     >>> cx.coords.replace_axes(xy, x, z)
     CartesianProduct(coordinates=(coordax.SizedAxis('z', size=4), coordax.SizedAxis('y', size=3)))
   """
+  # fmt: on
   if not to_replace.dims:
     raise ValueError(f'`to_replace` must have dimensions, got {to_replace}')
 
@@ -632,10 +642,11 @@ def from_xarray(
     coord_types: Sequence[type[Coordinate]] = (LabeledAxis, DummyAxis),
 ) -> Coordinate:
   # pylint: disable=line-too-long
-  """Convert the coordinates of an xarray.DataArray into a coordax.Coordinate.
+  # fmt: off
+  """Convert coordinates of an ``xarray.DataArray`` into ``coordax.Coordinate``.
 
   Args:
-    data_array: xarray.DataArray whose coordinates should be converted.
+    data_array: ``xarray.DataArray`` whose coordinates should be converted.
     coord_types: sequence of ``coordax.Coordinate`` subclasses with
       ``from_xarray`` methods defined. The first coordinate class that returns a
       coordinate object (indicating a match) will be used. By default,
@@ -643,8 +654,9 @@ def from_xarray(
       objects.
 
   Returns:
-    A coordax.Coordinate object representing the coordinates of the input
-    DataArray.
+    A coordax.Coordinate object, constructed from any of the indicated types
+    (plus ``CartesianProduct`` and ``Scalar``), representing the coordinates of
+    the input DataArray.
 
   Raises:
     ValueError: if no matching coordinate is found.
@@ -654,9 +666,13 @@ def from_xarray(
     >>> import xarray as xr
     >>> import numpy as np
     >>> da = xr.DataArray(np.zeros((2, 3)), dims=('x', 'y'), coords={'x': [1, 2]})
-    >>> cx.coords.from_xarray(da)  # doctest: +ELLIPSIS
+    >>> cx.coords.from_xarray(da)
     CartesianProduct(coordinates=(coordax.LabeledAxis('x', ticks=array([1, 2])), coordax.DummyAxis('y', size=3)))
+
+  See also:
+    :func:`coordax.from_xarray`
   """
+  # fmt: on
   dims = data_array.dims
   coords = []
 

coordax-0.2.2/coordax/experimental.py

@@ -0,0 +1,22 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Experimental parts of Coordax."""
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/jax-ml/jax/issues/7570
+# pylint: disable=g-multiple-import,useless-import-alias,g-importing-member,unused-import
+from coordax.ndarrays import (
+    NDArray as NDArray,
+    register_ndarray as register_ndarray,
+)

{coordax-0.2.1 → coordax-0.2.2}/coordax/fields.py

@@ -76,7 +76,7 @@ def _axes_attrs(field: Field) -> str:
 
 @utils.export
 def new_axis_name(field: Field, excluded_names: set[str] | None = None) -> str:
-  """Returns axis name that is not present in `field` or `excluded_names`.
+  """Returns axis name that is not present in ``field`` or ``excluded_names``.
 
   Args:
     field: The field to generate a new axis name for.
@@ -106,7 +106,7 @@ def new_axis_name(field: Field, excluded_names: set[str] | None = None) -> str:
 
 @utils.export
 def tmp_axis_name(field: Field, excluded_names: set[str] | None = None) -> str:
-  """Deprecated alias for `new_axis_name`."""
+  """Deprecated alias for ``new_axis_name``."""
   warnings.warn(
       'cx.tmp_axis_name() is deprecated, use cx.new_axis_name() instead',
       DeprecationWarning,
@@ -125,7 +125,7 @@ def cmap(
     vmap: Callable = jax.vmap,  # pylint: disable=g-bare-generic
 ) -> Callable[..., Any]:
   # fmt: off
-  """Vectorizes `fun` over coordinate dimensions of ``Field`` inputs.
+  """Vectorizes ``fun`` over coordinate dimensions of ``Field`` inputs.
 
   ``cmap`` is a "coordinate vectorizing map". It wraps an ordinary
   positional-axis-based function so that it accepts ``Field`` objects as input
@@ -244,7 +244,7 @@ def cpmap(
       coordax compatible with custom objects (e.g. neural net modules).
 
   Returns:
-    A function that applies `fun` to positional axes of the inputs while
+    A function that applies ``fun`` to positional axes of the inputs while
     vectorizing over all coordinate dimensions. The coordinate order of the
     inputs is preserved in the outputs.
 
@@ -264,8 +264,8 @@ def cpmap(
     Traceback (most recent call last):
     ...
     ValueError: 'same_as_input' for out_axes requires all NamedArray inputs with
-    named axes to have the same `named_axes`. Found multiple distinct
-    `named_axes`:
+    named axes to have the same named_axes. Found multiple distinct
+    named_axes on inputs:
     [{'x': 0, 'y': 1}, {'y': 0, 'x': 1}]
 
   See also:
@@ -314,8 +314,8 @@ def _cmap_with_doc(
     )
 
   docstr = (
-      f'Dimension-vectorized version of `{fun_name}`. Takes similar arguments'
-      f' as `{fun_name}` but accepts and returns Fields in place of arrays.'
+      f'Dimension-vectorized version of ``{fun_name}``. Takes similar arguments'
+      f' as ``{fun_name}`` but accepts and returns Fields in place of arrays.'
   )
   if fun_doc:
     docstr += f'\n\nOriginal documentation:\n\n{fun_doc}'
@@ -405,8 +405,8 @@ def _wrap_array_method(name):
   wrapped_func.__module__ = __name__
   wrapped_func.__doc__ = (
       'Name-vectorized version of array method'
-      f' `{name} <numpy.ndarray.{name}>`. Takes similar arguments as'
-      f' `{name} <numpy.ndarray.{name}>` but accepts and returns Fields'
+      f' ``{name} <numpy.ndarray.{name}>``. Takes similar arguments as'
+      f' ``{name} <numpy.ndarray.{name}>`` but accepts and returns Fields'
       ' in place of regular arrays.'
   )
   return wrapped_func
@@ -421,6 +421,7 @@ def _in_treescope_abbreviation_mode() -> bool:
 @jax.tree_util.register_pytree_node_class
 class Field:
   # pylint: disable=line-too-long
+  # fmt: off
   """An array with optional named dimensions and associated coordinates.
 
   Examples:
@@ -457,6 +458,7 @@ class Field:
     >>> field.coordinate
     CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.DummyAxis('y', size=3), coordax.DummyAxis(None, size=4)))
   """
+  # fmt on
 
   _named_array: named_axes_lib.NamedArray
   _axes: dict[str, Coordinate]
@@ -472,11 +474,11 @@ class Field:
     Args:
       data: the underlying data array.
       dims: optional tuple of dimension names, with the same length as
-        `data.ndim`. Strings indicate named axes, and may not be repeated.
-        `None` indicates positional axes. If `dims` is not provided, all axes
-        are positional.
+        ``data.ndim``. Strings indicate named axes, and may not be repeated.
+        ``None`` indicates positional axes. If ``dims`` is not provided, all
+        axes are positional.
       axes: optional mapping from dimension names to associated
-        `coordax.Coordinate` objects.
+        ``coordax.Coordinate`` objects.
 
     Examples:
       >>> import coordax as cx
@@ -645,8 +647,9 @@ class Field:
   def unwrap(self, *names: str | Coordinate) -> Array:
     """Extracts underlying data from a field without named dimensions.
 
-    This is effectively syntactic sugar for `assert field.named_dims == names`,
-    followed by returning `field.data`.
+    This is effectively syntactic sugar for ``assert field.named_dims ==
+    names``,
+    followed by returning ``field.data``.
 
     Args:
       *names: Names of dimensions to check against.
@@ -655,7 +658,7 @@ class Field:
       The underlying data array.
 
     Raises:
-      ValueError: If the field has named dimensions that do not match `names`.
+      ValueError: If the field has named dimensions that do not match ``names``.
 
     Examples:
       >>> import coordax as cx
@@ -678,7 +681,8 @@ class Field:
     return self.data
 
   def _validate_matching_coords(
-      self, dims_or_coords: Sequence[str | Coordinate | types.EllipsisType],
+      self,
+      dims_or_coords: Sequence[str | Coordinate | types.EllipsisType],
   ):
     """Validate that given coordinates are all found on this field."""
     axes = []
@@ -716,6 +720,8 @@ class Field:
 
       Untag by name:
 
+      >>> field.untag('x')
+      <Field dims=(None, 'y') shape=(2, 3) axes={} >
       >>> field.untag('y')
       <Field dims=('x', None) shape=(2, 3) axes={'x': SizedAxis} >
 
@@ -741,8 +747,8 @@ class Field:
     Args:
       *names: Names or coordinates to assign to the positional axes. The total
         number of dimensions corresponding to these objects must match the
-        number of positional axes in the field unless `...` is used to indicate
-        positions of untagged axes.
+        number of positional axes in the field unless ``...`` is used to
+        indicate positions of untagged axes.
 
     Returns:
       A new Field with the specified dimensions tagged.
@@ -754,18 +760,18 @@ class Field:
       >>> field.tag('x', 'y')
       <Field dims=('x', 'y') shape=(2, 3) axes={} >
 
-      Tagging with `Coordinate` objects adds them to the field:
+      Tagging with ``Coordinate`` objects adds them to the field:
 
       >>> x = cx.SizedAxis('x', 2)
       >>> field.tag(x, 'y')
       <Field dims=('x', 'y') shape=(2, 3) axes={'x': SizedAxis} >
 
-      `None` leaves a dimension untagged:
+      ``None`` leaves a dimension untagged:
 
       >>> field.tag('x', None)
       <Field dims=('x', None) shape=(2, 3) axes={} >
 
-      Ellipsis `...` can be indicate dimensions that should not be named:
+      Ellipsis ``...`` can be indicate dimensions that should not be named:
 
       >>> field.tag(..., 'y')
       <Field dims=(None, 'y') shape=(2, 3) axes={} >
@@ -775,7 +781,7 @@ class Field:
       >>> field.tag('x')  # doctest: +NORMALIZE_WHITESPACE
       Traceback (most recent call last):
       ...
-      ValueError: there must be exactly as many dimensions given to `tag` as
+      ValueError: there must be exactly as many dimensions given to ``tag`` as
       there are positional axes in the array, but got ('x',) for 2 positional
       axes.
 
@@ -790,7 +796,8 @@ class Field:
       CartesianProduct(axes={'x': SizedAxis, 'y': SizedAxis})
       >>> field = cx.Field(jnp.zeros((2, 3)))
       >>> field.tag(xy)
-      <Field dims=('x', 'y') shape=(2, 3) axes={'x': SizedAxis, 'y': SizedAxis} >
+      <Field dims=('x', 'y') shape=(2, 3) axes={'x': SizedAxis, 'y': SizedAxis}
+      >
 
     See also:
       :meth:`coordax.Field.untag`
@@ -812,12 +819,14 @@ class Field:
 
   # Note: Can't call this "transpose" like Xarray, to avoid conflicting with the
   # positional only ndarray method.
-  def order_as(self, *axis_order: str | Coordinate | types.EllipsisType) -> Field:
+  def order_as(
+      self, *axis_order: str | Coordinate | types.EllipsisType
+  ) -> Field:
     """Returns a field with the axes in the given order.
 
     Args:
       *axis_order: The desired order of axes, specified by name or coordinate.
-        `...` may be used once, to indicate all other dimensions in order of
+        ``...`` may be used once, to indicate all other dimensions in order of
         appearance on this array.
 
     Returns:
@@ -839,13 +848,13 @@ class Field:
     return result
 
   def broadcast_like(self, other: Self | Coordinate) -> Self:
-    """Returns a field broadcasted like `other`.
+    """Returns a field broadcasted like ``other``.
 
     Args:
       other: The field or coordinate to broadcast to.
 
     Returns:
-      A new Field broadcasted to match `other`.
+      A new Field broadcasted to match ``other``.
 
     Examples:
       >>> import coordax as cx
@@ -1010,7 +1019,7 @@ class Field:
 def field(array: ArrayLike, *names: str | Coordinate | None) -> Field:
   """Wraps a positional array as a ``Field``.
 
-  `cx.field(array, *names)` is a shortcut for `cx.Field(array).tag(*names)`.
+  ``cx.field(array, *names)`` is a shortcut for ``cx.Field(array).tag(*names)``.
 
   Args:
     array: the array to wrap.
@@ -1037,7 +1046,7 @@ def field(array: ArrayLike, *names: str | Coordinate | None) -> Field:
 
 @utils.export
 def wrap(array: ArrayLike, *names: str | Coordinate | None) -> Field:
-  """Deprecated alias for `cx.field`."""
+  """Deprecated alias for ``cx.field``."""
   warnings.warn(
       'cx.wrap() is deprecated, use cx.field() instead',
       DeprecationWarning,
@@ -1048,7 +1057,7 @@ def wrap(array: ArrayLike, *names: str | Coordinate | None) -> Field:
 
 @utils.export
 def wrap_like(array: ArrayLike, other: Field) -> Field:
-  """Wraps `array` with the same coordinates as `other`."""
+  """Wraps ``array`` with the same coordinates as ``other``."""
   warnings.warn(
       'cx.wrap_like() is deprecated, use cx.field(array, other.coordinate) '
       'instead of cx.wrap_like(array, other)',
@@ -1068,6 +1077,7 @@ def from_xarray(
     coord_types: Sequence[type[Coordinate]] = (LabeledAxis, DummyAxis),
 ) -> Field:
   # pylint: disable=g-import-not-at-top,line-too-long
+  # fmt: off
   """Create a coordax.Field from an xarray.DataArray.
 
   Args:
@@ -1092,6 +1102,7 @@ def from_xarray(
   See also:
     :func:`coordax.coords.from_xarray`
   """
+  # fmt: on
   data = data_array.data
   coord = coordinate_systems.from_xarray(data_array, coord_types)
   return field(data, coord)
@@ -1099,19 +1110,19 @@ def from_xarray(
 
 @utils.export
 def is_field(value) -> TypeGuard[Field]:
-  """Returns True if `value` is of type `Field`."""
+  """Returns True if ``value`` is of type ``Field``."""
   return isinstance(value, Field)
 
 
 @utils.export
 def shape_struct_field(*axes: Coordinate) -> Field:
-  """Returns a Field with `axes` and a ShapeDtypeStruct in place of data.
+  """Returns a Field with ``axes`` and a ShapeDtypeStruct in place of data.
 
   Args:
     *axes: The coordinates to use for the field.
 
   Returns:
-    A Field with `ShapeDtypeStruct` data and the given axes.
+    A Field with ``ShapeDtypeStruct`` data and the given axes.
 
   Examples:
     >>> import coordax as cx
@@ -1142,12 +1153,12 @@ def get_coordinate(
     field: coordax.Field from which the coordinate will be extracted.
     missing_axes: controls how axes without coorinates are handled. Options are:
 
-      * ``'dummy'``: uses DummyAxis for dimensions without a coordinate.
+      * ``'dummy'``: uses ``DummyAxis`` for dimensions without a coordinate.
       * ``'skip'``: ignores dimensions without a coordinate.
       * ``'error'``: raises if dimensions without a coordinate are present.
 
   Returns:
-    Coordinate associated with the `field`.
+    Coordinate associated with the ``field``.
 
   Examples:
     >>> import coordax as cx

{coordax-0.2.1 → coordax-0.2.2}/coordax/named_axes.py

@@ -15,8 +15,8 @@
 
 This module is intended to be nearly a drop-in replacement for Penzai's
 NamedArray, but with an alternative, simpler implementation. Dimensions are
-specified by a tuple, where each element is either a string or `None`. `None` is
-used to indicate strictly positional dimensions.
+specified by a tuple, where each element is either a string or ``None``.
+``None`` is used to indicate strictly positional dimensions.
 
 Some code and documentation is adapted from penzai.core.named_axes.
 """
@@ -49,7 +49,7 @@ def attrs_summary_type(
     named_array: NamedArray,
     inspect_data: bool,
 ) -> tuple[str, str, str]:
-  """Returns a summary of a `named_array` and its data type."""
+  """Returns a summary of a ``named_array`` and its data type."""
   if isinstance(named_array.data, jax.Array) and not isinstance(
       named_array.data, jax.core.Tracer
   ):
@@ -111,8 +111,8 @@ class NamedArrayAdapter(ndarray_adapters.NDArrayAdapter['NamedArray']):
         bad_names = set(mask.dims) - set(array.dims)
         if bad_names:
           raise ValueError(
-              'Valid mask must be broadcastable to the shape of `array`, but it'
-              f' had extra axis names {bad_names}'
+              'Valid mask must be broadcastable to the shape of ``array``, but '
+              f'it had extra axis names {bad_names}'
           )
         mask = mask.order_as(*(d for d in array.dims if d in mask.dims))
         mask_data = mask.data
@@ -240,8 +240,8 @@ def _normalize_out_axes(
         named_axes_list = [arr.named_axes for arr in named_arrays_with_axes]
         raise ValueError(
             "'same_as_input' for out_axes requires all NamedArray inputs with"
-            ' named axes to have the same `named_axes`. Found multiple'
-            f' distinct `named_axes`:\n{named_axes_list}'
+            ' named axes to have the same named_axes. Found multiple'
+            f' distinct named_axes on inputs:\n{named_axes_list}'
         )
       [unique_named_axes] = unique_named_axes
       return dict(unique_named_axes)
@@ -304,22 +304,22 @@ def nmap(
   """Automatically vectorizes ``fun`` over named dimensions.
 
   ``nmap`` is a "named dimension vectorizing map". It wraps an ordinary
-  positional-axis-based function so that it accepts NamedArrays as input and
-  produces NamedArrays as output, and vectorizes over all of named dimensions,
+  positional-axis-based function so that it accepts ``NamedArray``s as input and
+  produces ``NamedArray``s as output, and vectorizes over all of named dimensions,
   calling the original function with positionally-indexed slices corresponding
-  to each argument's `positional_shape`.
+  to each argument's ``positional_shape``.
 
-  Unlike `jax.vmap`, the axes to vectorize over are inferred
+  Unlike ``jax.vmap``, the axes to vectorize over are inferred
   automatically from the named dimensions in the NamedArray inputs, rather
   than being specified as part of the mapping transformation. Specifically, each
   dimension name that appears in any of the arguments is vectorized over jointly
   across all arguments that include that dimension, and is then included as a
   named dimension in the output. To make an axis visible to ``fun``, you can
-  call `untag` on the argument and pass the axis name(s) of interest; ``fun``
+  call ``untag`` on the argument and pass the axis name(s) of interest; ``fun``
   will then see those axes as positional axes instead of mapping over them.
 
-  `untag` and ``nmap`` are together the primary ways to apply individual
-  operations to axes of a NamedArray. `tag` can then be used on the result to
+  ``untag`` and ``nmap`` are together the primary ways to apply individual
+  operations to axes of a NamedArray. ``tag`` can then be used on the result to
   re-bind names to positional axes.
 
   Within ``fun``, any mapped-over axes will be accessible using standard JAX
@@ -347,8 +347,8 @@ def nmap(
 
   Returns:
     An automatically-vectorized version of ``fun``, which can optionally be
-    called with NamedArrays instead of ordinary arrays, and which will always
-    return NamedArrays for each of its output leaves. Any argument (or PyTree
+    called with ``NamedArray``s instead of ordinary arrays, and which will always
+    return ``NamedArray``s for each of its output leaves. Any argument (or PyTree
     leaf of an argument) that is a NamedArray will have its named dimensions
     vectorized over; ``fun`` will then be called with batch tracers
     corresponding to slices of the input array that are shaped like
@@ -452,8 +452,8 @@ def _nmap_with_doc(
     return jax.tree.map(wrap_output, result, is_leaf=is_array)
 
   docstr = (
-      f'Dimension-vectorized version of `{fun_name}`. Takes similar arguments'
-      f' as `{fun_name}` but accepts and returns NamedArray objects in place '
+      f'Dimension-vectorized version of ``{fun_name}``. Takes similar arguments'
+      f' as ``{fun_name}`` but accepts and returns NamedArray objects in place '
       'of arrays.'
   )
   if fun_doc:
@@ -536,8 +536,8 @@ def _wrap_array_method(
   wrapped_func.__module__ = __name__
   wrapped_func.__doc__ = (
       'Name-vectorized version of array method'
-      f' `{name} <numpy.ndarray.{name}>`. Takes similar arguments as'
-      f' `{name} <numpy.ndarray.{name}>` but accepts and returns NamedArray'
+      f' ``{name} <numpy.ndarray.{name}>``. Takes similar arguments as'
+      f' ``{name} <numpy.ndarray.{name}>`` but accepts and returns NamedArray'
       ' objects in place of regular arrays.'
   )
   return wrapped_func
@@ -560,7 +560,7 @@ class _ShapedLeaf:
 
 
 def _tmp_axis_name(x: NamedArray, excluded_names: set[str]) -> str:
-  """Returns axis name that is not present in `x` or `excluded_names`."""
+  """Returns axis name that is not present in ``x`` or ``excluded_names``."""
   for i in range(x.ndim):
     name = f'tmp_axis_{i}'
     if name not in excluded_names and name not in x.named_dims:
@@ -582,9 +582,9 @@ class NamedArray:
 
   Attributes:
     data: the underlying data array.
-    dims: tuple of dimension names, with the same length as `data.ndim`. Strings
-      indicate named axes, and may not be repeated. `None` indicates positional
-      axes.
+    dims: tuple of dimension names, with the same length as ``data.ndim``.
+      Strings indicate named axes, and may not be repeated. ``None`` indicates
+      positional axes.
   """
 
   _data: Array
@@ -600,9 +600,9 @@ class NamedArray:
     Arguments:
       data: the underlying data array.
       dims: optional tuple of dimension names, with the same length as
-        `data.ndim`. Strings indicate named axes, and may not be repeated.
-        `None` indicates positional axes. If `dims` is not provided, all axes
-        are positional.
+        ``data.ndim``. Strings indicate named axes, and may not be repeated.
+        ``None`` indicates positional axes. If ``dims`` is not provided, all
+        axes are positional.
     """
     data = ndarrays.to_array(data)
     if dims is None:
@@ -804,8 +804,8 @@ class NamedArray:
   def untag(self, *dims: str) -> Self:
     """Removes the requested dimension names.
 
-    `untag` can only be called on a `NamedArray` that does not have any
-    positional axes. It produces a new `NamedArray` where the axes with the
+    ``untag`` can only be called on a ``NamedArray`` that does not have any
+    positional axes. It produces a new ``NamedArray`` where the axes with the
     requested dimension names are now treated as positional instead.
 
     Args:
@@ -846,12 +846,12 @@ class NamedArray:
   def order_as(self, *dims: str | types.EllipsisType) -> Self:
     """Reorder the dimensions of an array.
 
-    All dimensions must be named. Use `tag` first to name any positional axes.
+    All dimensions must be named. Use ``tag`` first to name any positional axes.
 
     Args:
       *dims: dimension names that appear on this array, in the desired order on
-        the result. `...` may be used once, to indicate all other dimensions in
-        order of appearance on this array.
+        the result. ``...`` may be used once, to indicate all other dimensions
+        in order of appearance on this array.
 
     Returns:
       Array with transposed data and reordered dimensions, as indicated.

{coordax-0.2.1 → coordax-0.2.2}/coordax/named_axes_test.py

@@ -110,7 +110,7 @@ class NamedAxesTest(parameterized.TestCase):
     with self.assertRaisesWithLiteralMatch(
         TypeError,
         'data must be a np.ndarray, jax.Array or a duck-typed array registered'
-        ' with coordax.register_ndarray(), got dict: {}',
+        ' with coordax.experimental.register_ndarray(), got dict: {}',
     ):
       named_axes.NamedArray({})
     with self.assertRaisesRegex(
@@ -784,9 +784,9 @@ class NamedAxesTest(parameterized.TestCase):
         ValueError,
         re.escape(
             "'same_as_input' for out_axes requires all NamedArray inputs with"
-            ' named axes to have the same `named_axes`. Found multiple'
-            " distinct `named_axes`:\n[{'x': 0, 'y': 1, 'z': 2}, {'z': 0,"
-            " 'y': 1, 'x': 2}]"
+            ' named axes to have the same named_axes. Found multiple'
+            " distinct named_axes on inputs:\n"
+            "[{'x': 0, 'y': 1, 'z': 2}, {'z': 0, 'y': 1, 'x': 2}]"
         ),
     ):
       named_axes.nmap(lambda x, y: x, out_axes='same_as_input')(array1, array2)
@@ -796,8 +796,8 @@ class NamedAxesTest(parameterized.TestCase):
     with self.assertRaisesWithLiteralMatch(
         ValueError,
         "'same_as_input' for out_axes requires all NamedArray inputs with"
-        ' named axes to have the same `named_axes`. Found multiple'
-        " distinct `named_axes`:\n[{'x': 0}, {'x': 0, 'y': 1}]"
+        ' named axes to have the same named_axes. Found multiple'
+        " distinct named_axes on inputs:\n[{'x': 0}, {'x': 0, 'y': 1}]"
     ):
       named_axes.nmap(lambda x, y: x, out_axes='same_as_input')(array1, array2)
 

{coordax-0.2.1 → coordax-0.2.2}/coordax/ndarrays.py

@@ -16,8 +16,10 @@
 from __future__ import annotations
 
 import abc
+import functools
 from typing import Any, Callable, Self, TypeVar
 
+from coordax import utils
 import jax
 import numpy as np
 
@@ -28,10 +30,15 @@ except ImportError:
 
 
 # TODO(shoyer): should this be a protocol?
+@functools.partial(utils.export, module='coordax.experimental')
 class NDArray(abc.ABC):
   """Abstract base class for non-JAX arrays that can be used with Coordax.
 
-  To register a new NDArray, use `coordax.register_ndarray()`.
+  THIS FEATURE IS EXPERIMENTAL, AND THE DETAILS ARE SUBJECT TO CHANGE. We
+  currently use it only to support jax-datetime, but expect it might be useful
+  in the future for other array types, e.g., with arbitrary units.
+
+  To register a new NDArray, use :func:`coordax.experimental.register_ndarray`.
 
   NDArray instances are also expected to be registered as JAX pytree nodes:
   https://docs.jax.dev/en/latest/pytrees.html#extending-pytrees
@@ -45,12 +52,12 @@ class NDArray(abc.ABC):
   @property
   @abc.abstractmethod
   def size(self) -> int:
-    """Size of this array, typically `math.prod(self.shape)`."""
+    """Size of this array, typically ``math.prod(self.shape)``."""
 
   @property
   @abc.abstractmethod
   def ndim(self) -> int:
-    """Number of dimensions in this array, typcially `len(self.shape)`."""
+    """Number of dimensions in this array, typcially ``len(self.shape)``."""
 
   @abc.abstractmethod
   def transpose(self, axes: tuple[int, ...]) -> Self:
@@ -58,7 +65,7 @@ class NDArray(abc.ABC):
 
   # Note: __getitem__ is not yet used by Coordax, but we will likely want it in
   # the near the future.
-  # TODO(shoyer): Figure out the precise type for `value`.
+  # TODO(shoyer): Figure out the precise type for ``value``.
   @abc.abstractmethod
   def __getitem__(self, value) -> Self:
     """Index this array, returning a new array."""
@@ -72,7 +79,7 @@ Array = np.ndarray | jax.Array | NDArray
 
 
 def to_array(data: ArrayLike) -> Array:
-  """Returns an NDArray compatible with JAX and Coordax.
+  """Returns an unlabeled NDArray compatible with JAX and Coordax.
 
   Args:
     data: a scalar, NumPy array, JAX array, or registered NDArray.
@@ -93,7 +100,8 @@ def to_array(data: ArrayLike) -> Array:
   if not isinstance(data, (np.ndarray, jax.Array, NDArray)):
     raise TypeError(
         'data must be a np.ndarray, jax.Array or a duck-typed array registered '
-        f'with coordax.register_ndarray(), got {type(data).__name__}: {data}'
+        'with coordax.experimental.register_ndarray(), got '
+        f'{type(data).__name__}: {data}'
     )
 
   return data
@@ -119,13 +127,19 @@ _FROM_NUMPY_FUNCS: list[
 ] = []
 
 
+@functools.partial(utils.export, module='coordax.experimental')
 def register_ndarray(
     array_type: type[T],
     is_matching_numpy_array: Callable[[np.ndarray], bool],
     to_numpy: Callable[[T], np.ndarray],
     from_numpy: Callable[[np.ndarray], T],
 ) -> type[T]:
-  """Registers a new duck-typed array type corresponding to a numpy dtype."""
+  # pylint: disable=g-doc-args
+  """Registers a new duck-typed array type corresponding to a numpy dtype.
+
+  See also:
+    :class:`coordax.experimental.NDArray`
+  """
   NDArray.register(array_type)
   _TO_NUMPY_FUNCS.append((array_type, to_numpy))
   _FROM_NUMPY_FUNCS.append((is_matching_numpy_array, from_numpy))

{coordax-0.2.1 → coordax-0.2.2}/coordax/testing.py

@@ -20,7 +20,7 @@ def assert_field_properties(
     rtol: float | None = 1e-5,
     atol: float | None = 1e-5,
 ):
-  """Asserts that a Field has expected properties."""
+  """Asserts that a ``Field`` has expected properties."""
   if data is not None:
     if atol is None and rtol is None:
       chex.assert_trees_all_equal(actual.data, data)
@@ -48,7 +48,7 @@ def assert_fields_allclose(
     rtol: float = 1e-5,
     atol: float = 1e-5,
 ):
-  """Asserts that two Fields are close and have matching coordinates."""
+  """Asserts that two ``Field``s are close and have matching coordinates."""
   assert_field_properties(
       actual=actual,
       data=desired.data,
@@ -63,7 +63,7 @@ def assert_fields_allclose(
 
 
 def assert_fields_equal(actual: coordax.Field, desired: coordax.Field):
-  """Asserts that two Fields are equal and have matching coordinates."""
+  """Asserts that two ``Field``s are equal and have matching coordinates."""
   assert_field_properties(
       actual=actual,
       data=desired.data,

{coordax-0.2.1 → coordax-0.2.2}/coordax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.1
+Version: 0.2.2
 Summary: Coordinate axes for scientific computing in JAX
 Author-email: Google LLC <noreply@google.com>
 License-Expression: Apache-2.0

{coordax-0.2.1 → coordax-0.2.2}/coordax.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ coordax/coordinate_systems.py
 coordax/coordinate_systems_test.py
 coordax/coords.py
 coordax/doctest_test.py
+coordax/experimental.py
 coordax/fields.py
 coordax/fields_test.py
 coordax/jax_datetime_integration_test.py

{coordax-0.2.1 → coordax-0.2.2}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["coordax"]
 
 [project]
 name = "coordax"
-version = "0.2.1"  # keep sync with __init__.py
+version = "0.2.2"  # keep sync with __init__.py
 description = "Coordinate axes for scientific computing in JAX"
 authors = [
     {name = "Google LLC", email = "noreply@google.com"},