check_shapes User Guide

A library for annotating and checking the shapes of tensors.

• Documentation for most recent release

• Documentation for current development version

• User Guide
• Benchmark results
• API reference

The main entry point is check_shapes().

Example:

import numpy as np

from check_shapes import check_shapes

@check_shapes(
    "features: [batch..., n_features]",
    "weights: [n_features]",
    "return: [batch...]",
)
def linear_model(features: ndarray, weights: ndarray) -> ndarray:
    return np.einsum("...i,i -> ...", features, weights)

Main features include:

• Supports NumPy, TensorFlow, JAX and PyTorch out-of-the-box, and easy to extend to other frameworks.

• Checking the shapes of function arguments and return types, using a decorator.

• Checking the shapes of local values, using a function.

• Constant- and variable size dimensions.

• Constant- and variable rank tensors.

• Broadcasting.

• Conditional shapes.

• Reuse of shape specifications, including in class inheritance.

• Automatic rewrite of docstrings to include shape information.

Installation

check_shapes can be installed with pip as usual:

pip install check_shapes

Speed, and interactions with tf.function

Shape checking has some performance impact. For estimates of this impact see our Benchmark results page. If the overhead is unacceptable shape checking can be disabled. Shape checking can be set to one of three different states:

• ENABLED. Shapes are checked wherever they can be.

• EAGER_MODE_ONLY. Shapes are not checked within anything wrapped in tf.function().

• DISABLED. Shapes are never checked.

The state can be set with set_enable_check_shapes():

set_enable_check_shapes(ShapeCheckingState.DISABLED)

Alternatively you can use disable_check_shapes() to disable shape checking in smaller scopes:

with disable_check_shapes():
    performance_sensitive_function()

Beware that any function declared while shape checking is disabled, will continue not to check shapes, even if shape checking is otherwise enabled again.

The default state is EAGER_MODE_ONLY; which is appropriate for smaller project, experiments, and notebooks. Write and debug your code in eager mode, and add tf.function() when you believe your code is correct and you want it to run fast. For larger project you probably want to modify this setting. In particular you may want to enable all shape checks in your unit tests. If you use pytest you can do this by updating your root conftest.py with:

@pytest.fixture(autouse=True)
def enable_shape_checks() -> Iterable[None]:
    old_enable = get_enable_check_shapes()
    old_rewrite_docstrings = get_rewrite_docstrings()
    old_function_call_precompute = get_enable_function_call_precompute()
    set_enable_check_shapes(ShapeCheckingState.ENABLED)
    set_rewrite_docstrings(DocstringFormat.SPHINX)
    set_enable_function_call_precompute(True)
    yield
    set_enable_function_call_precompute(old_function_call_precompute)
    set_rewrite_docstrings(old_rewrite_docstrings)
    set_enable_check_shapes(old_enable)

If shape checking is set to ENABLED and your code is wrapped in tf.function() shape checks are performed while tracing graphs, but not compiled into the actual graphs. This is considered a feature as that means that check_shapes() doesn’t impact the execution speed of your functions after they have been compiled.

Best-effort checking

This library will perform shape checks on a best-effort basis. Many things can prevent this library from being able to check shapes. For example:

• Unknown shapes. Sometimes the library is not able to determine the shape of an object, and thus cannot check that object. For example Optional arguments with value None cannot be checked, and compiled TensorFlow code can have variables with an unknown shape.

• Use of variable-rank dimensions (see below). In general we cannot infer the size of variable-rank dimensions if there are multiple variable-rank specifications within the same shape specification (e.g. cov: [m..., n...]). This library will try to learn the size of these variable-rank dimensions from neighbouring shape specifications, but this is not always possible. Use of broadcast with variable-rank dimensions makes it even harder to infer these values.

Check specification

The shapes to check are specified by the arguments to check_shapes(). Each argument is a string of the format:

<argument specifier> ":" <shape specifier> ["if" <condition>] ["#" <note>]

Argument specification

The <argument specifier> must start with either the name of an argument to the decorated function, or the special name return. The value return refers to the value returned by the function.

The <argument specifier> can then be modified to refer to elements of the object in several ways:

• Use .<name> to refer to an attribute of an object:

  
  @dataclass
  class Statistics:
      mean: ndarray
      std: ndarray
  
  @check_shapes(
      "data: [n_rows, n_columns]",
      "return.mean: [n_columns]",
      "return.std: [n_columns]",
  )
  def compute_statistics(data: ndarray) -> Statistics:
      return Statistics(np.mean(data, axis=0), np.std(data, axis=0))
  
  

• Use [<index>] to refer to a specific element of a sequence. This is particularly useful if your function returns a tuple of values:

  
  @check_shapes(
      "data: [n_rows, n_columns]",
      "return[0]: [n_columns]",
      "return[1]: [n_columns]",
  )
  def compute_mean_and_std(data: ndarray) -> Tuple[ndarray, ndarray]:
      return np.mean(data, axis=0), np.std(data, axis=0)
  
  

• Use [all] to select all elements of a collection:

  
  @check_shapes(
      "data[all]: [., n_columns]",
      "return: [., n_columns]",
  )
  def concat_rows(data: Sequence[ndarray]) -> ndarray:
      return np.concatenate(data, axis=0)
  
  concat_rows(
      [
          np.ones((1, 3)),
          np.ones((4, 3)),
      ]
  )
  
  

• Use .keys() to select all keys of a mapping:

  
  @check_shapes(
      "data.keys(): [.]",
      "return: []",
  )
  def sum_key_lengths(data: Mapping[Tuple[int, ...], str]) -> int:
      return sum(len(k) for k in data)
  
  sum_key_lengths(
      {
          (3,): "foo",
          (1, 2): "bar",
      }
  )
  
  

• Use .values() to select all values of a mapping:

  
  @check_shapes(
      "data.values(): [., n_columns]",
      "return: [., n_columns]",
  )
  def concat_rows(data: Mapping[str, ndarray]) -> ndarray:
      return np.concatenate(list(data.values()), axis=0)
  
  concat_rows(
      {
          "foo": np.ones((1, 3)),
          "bar": np.ones((4, 3)),
      }
  )
  
  

Note

We do not support looking up a specific key or value in a dict.

If the argument, or any of the looked-up values, are None the check is skipped. This is useful for optional values:

@check_shapes(
    "x1: [n_rows_1, n_inputs]",
    "x2: [n_rows_2, n_inputs]",
    "return: [n_rows_1, n_rows_2]",
)
def squared_exponential_kernel(
    variance: float, x1: ndarray, x2: Optional[ndarray] = None
) -> ndarray:
    if x2 is None:
        x2 = x1
    cov: ndarray = variance * np.exp(
        -0.5 * np.sum((x1[:, None, :] - x2[None, :, :]) ** 2, axis=2)
    )
    return cov

squared_exponential_kernel(1.0, np.ones((3, 2)), np.ones((4, 2)))
squared_exponential_kernel(3.2, np.ones((3, 2)))

Shape specification

Shapes are specified by the syntax:

"[" <dimension specifier 1> "," <dimension specifer 2> "," ... "," <dimension specifier n> "]"

where <dimension specifier i> is one of:

• <integer>, to require that dimension to have that exact size:

  
  @check_shapes(
      "v1: [2]",
      "v2: [2]",
  )
  def vector_2d_distance(v1: ndarray, v2: ndarray) -> float:
      return float(np.sqrt(np.sum((v1 - v2) ** 2)))
  
  

• <name>, to bind that dimension to a variable. Dimensions bound to the same variable must have the same size, though that size can be anything:

  
  @check_shapes(
      "v1: [d]",
      "v2: [d]",
  )
  def vector_distance(v1: ndarray, v2: ndarray) -> float:
      return float(np.sqrt(np.sum((v1 - v2) ** 2)))
  
  

• None or . to allow exactly one single dimension without constraints:

  
  @check_shapes(
      "v: [None]",
  )
  def vector_length(v: ndarray) -> float:
      return float(np.sqrt(np.sum(v ** 2)))
  
  

  or:

  
  @check_shapes(
      "v: [.]",
  )
  def vector_length(v: ndarray) -> float:
      return float(np.sqrt(np.sum(v ** 2)))
  
  

• *<name> or <name>..., to bind any number of dimensions to a variable. Again, multiple uses of the same variable name must match the same dimension sizes:

  
  @check_shapes(
      "x: [*batch, n_columns]",
      "return: [*batch]",
  )
  def batch_mean(x: ndarray) -> ndarray:
      mean: ndarray = np.mean(x, axis=-1)
      return mean
  
  

  or:

  
  @check_shapes(
      "x: [batch..., n_columns]",
      "return: [batch...]",
  )
  def batch_mean(x: ndarray) -> ndarray:
      mean: ndarray = np.mean(x, axis=-1)
      return mean
  
  

• * or ..., to allow any number of dimensions without constraints:

  
  @check_shapes(
      "x: [*]",
  )
  def rank(x: ndarray) -> int:
      return len(x.shape)
  
  

  or:

  
  @check_shapes(
      "x: [...]",
  )
  def rank(x: ndarray) -> int:
      return len(x.shape)
  
  

A scalar shape is specified by []:

@check_shapes(
    "x: [...]",
    "return: []",
)
def mean(x: ndarray) -> ndarray:
    mean: ndarray = np.sum(x) / x.size
    return mean

Any of the above can be prefixed with the keyword broadcast to allow any value that broadcasts to the specification. For example:

@check_shapes(
    "a: [broadcast batch...]",
    "b: [broadcast batch...]",
    "return: [batch...]",
)
def add(a: ndarray, b: ndarray) -> ndarray:
    return a + b

Specifically, to mark a dimension as broadcast means:

• If the specification is that the dimension should have size n, then the actual dimension must have value 1 or n.

• If all leading dimension specifications are also marked broadcast, then the actual shape is allowed to be shorter than the specification — the dimension is allowed to be missing.

Condition specification

You can use the optional if <condition> syntax to conditionally evaluate shapes. If an if <condition> is used, the specification is only appplied if <condition> evaluates to True. This is useful if shapes depend on other input parameters. Valid conditions are:

• <argument specifier>, with the same syntax and rules as above, except that constructions that evaluates to multiple elements are disallowed. Uses the bool built-in to convert the value of the argument to a bool:

  
  @check_shapes(
      "a: [broadcast batch...] if check_a",
      "b: [broadcast batch...] if check_b",
      "return: [batch...]",
  )
  def add(a: ndarray, b: ndarray, check_a: bool = True, check_b: bool = True) -> ndarray:
      return a + b
  
  add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)
  
  

• <argument specifier> is None, and <argument specifier> is not None, with the usual rules for an <argument specifier>, to test whether an argument is, or is not, None. We currently only allow tests against None, not general Python equality tests:

  
  @check_shapes(
      "a: [n_a]",
      "b: [n_b]",
      "return: [n_a, n_a] if b is None",
      "return: [n_a, n_b] if b is not None",
  )
  def square(a: ndarray, b: Optional[ndarray] = None) -> ndarray:
      if b is None:
          b = a
      result: ndarray = a[:, None] * b[None, :]
      return result
  
  square(np.ones((3,)))
  square(np.ones((3,)), np.ones((4,)))
  
  

• <left> or <right>, evaluates to True if any of <left> or <right> evaluates to True and to False otherwise:

  
  @check_shapes(
      "a: [broadcast batch...] if check_all or check_a",
      "b: [broadcast batch...] if check_all or check_b",
      "return: [batch...]",
  )
  def add(
      a: ndarray,
      b: ndarray,
      check_all: bool = False,
      check_a: bool = True,
      check_b: bool = True,
  ) -> ndarray:
      return a + b
  
  add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)
  
  

• <left> and <right>, evaluates to False if any of <left> or <right> evaluates to False and to True otherwise:

  
  @check_shapes(
      "a: [broadcast batch...] if enable_checks and check_a",
      "b: [broadcast batch...] if enable_checks and check_b",
      "return: [batch...]",
  )
  def add(
      a: ndarray,
      b: ndarray,
      enable_checks: bool = True,
      check_a: bool = True,
      check_b: bool = True,
  ) -> ndarray:
      return a + b
  
  add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)
  
  

• not <right>, evaluates to the opposite value of <right>:

  
  @check_shapes(
      "a: [broadcast batch...] if not disable_checks",
      "b: [broadcast batch...] if not disable_checks",
      "return: [batch...]",
  )
  def add(a: ndarray, b: ndarray, disable_checks: bool = False) -> ndarray:
      return a + b
  
  add(np.ones((3, 1)), np.ones((1, 4)))
  
  

• (<exp>), uses parenthesis to change operator precedence, as usual.

Conditions can be composed to apply different specs, depending on function arguments:

@check_shapes(
    "a: [j] if a_vector",
    "a: [i, j] if (not a_vector)",
    "b: [j] if b_vector",
    "b: [j, k] if (not b_vector)",
    "return: [1, 1] if a_vector and b_vector",
    "return: [1, k] if a_vector and (not b_vector)",
    "return: [i, 1] if (not a_vector) and b_vector",
    "return: [i, k] if (not a_vector) and (not b_vector)",
)
def multiply(a: ndarray, b: ndarray, a_vector: bool, b_vector: bool) -> ndarray:
    if a_vector:
        a = a[None, :]
    if b_vector:
        b = b[:, None]

    return a @ b

multiply(np.ones((4,)), np.ones((4, 5)), a_vector=True, b_vector=False)

Note

All specifications with either no if syntax or a <condition> that evaluates to True will be applied. It is possible for multiple specifications to apply to the same value.

Note specification

You can add notes to your specifications using a # followed by the note. These notes will be appended to relevant error messages and appear in rewritten docstrings. You can add notes in two places:

• On a single line by itself, to add a note to the entire function:

  
  @check_shapes(
      "features: [batch..., n_features]",
      "# linear_model currently only supports a single output.",
      "weights: [n_features]",
      "return: [batch...]",
  )
  def linear_model(features: ndarray, weights: ndarray) -> ndarray:
      prediction: ndarray = np.einsum("...i,i -> ...", features, weights)
      return prediction
  
  

• After the specification of a single argument, to add a note to that argument only:

  
  @check_shapes(
      "features: [batch..., n_features]",
      "weights: [n_features] # linear_model currently only supports a single output.",
      "return: [batch...]",
  )
  def linear_model(features: ndarray, weights: ndarray) -> ndarray:
      prediction: ndarray = np.einsum("...i,i -> ...", features, weights)
      return prediction
  
  

Shape reuse

Just like with other code it is useful to be able to specify a shape in one place and reuse the specification. In particular this ensures that your code keep having internally consistent shapes, even if it is refactored.

Class inheritance

If you have a class hiererchy, you probably want to ensure that derived classes handle tensors with the same shapes as the base classes. You can use the inherit_check_shapes() decorator to inherit shapes from overridden methods:

class Model(ABC):
    @abstractmethod
    @check_shapes(
        "features: [batch..., n_features]",
        "return: [batch...]",
    )
    def predict(self, features: ndarray) -> ndarray:
        pass

class LinearModel(Model):
    @check_shapes(
        "weights: [n_features]",
    )
    def __init__(self, weights: ndarray) -> None:
        self._weights = weights

    @inherit_check_shapes
    def predict(self, features: ndarray) -> ndarray:
        prediction: ndarray = np.einsum("...i,i -> ...", features, self._weights)
        return prediction

Functional programming

If you prefer functional- over object oriented programming, you may have functions that you require to handle the same shapes. To do this, remember that in Python a decorator is just a function, and functions are objects that can be stored:

check_metric_shapes = check_shapes(
    "actual: [n_rows, n_labels]",
    "predicted: [n_rows, n_labels]",
    "return: []",
)

@check_metric_shapes
def rmse(actual: ndarray, predicted: ndarray) -> float:
    return float(np.mean(np.sqrt(np.mean((predicted - actual) ** 2, axis=-1))))

@check_metric_shapes
def mape(actual: ndarray, predicted: ndarray) -> float:
    return float(np.mean(np.abs((predicted - actual) / actual)))

Other reuse of shapes

You can use get_check_shapes() to get, and reuse, the shape definitions from a previously declared function. This is particularly useful to ensure fakes in tests use the same shapes as the production implementation:

class Model(ABC):
    @abstractmethod
    @check_shapes(
        "features: [batch..., n_features]",
        "return: [batch...]",
    )
    def predict(self, features: ndarray) -> ndarray:
        pass

@check_shapes(
    "test_features: [n_rows, n_features]",
    "test_labels: [n_rows]",
)
def evaluate_model(model: Model, test_features: ndarray, test_labels: ndarray) -> float:
    prediction = model.predict(test_features)
    return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

def test_evaluate_model() -> None:
    fake_features = np.ones((10, 3))
    fake_labels = np.ones((10,))
    fake_predictions = np.ones((10,))

    @get_check_shapes(Model.predict)
    def fake_predict(features: ndarray) -> ndarray:
        assert features is fake_features
        return fake_predictions

    fake_model = MagicMock(spec=Model, predict=fake_predict)

    assert pytest.approx(0.0) == evaluate_model(fake_model, fake_features, fake_labels)

Checking intermediate results

You can use the function check_shape() to check the shape of an intermediate result. This function will use the same namespace as the immediately surrounding check_shapes() decorator, and should only be called within functions that has such a decorator. For example:

@check_shapes(
    "weights: [n_features, n_labels]",
    "test_features: [n_rows, n_features]",
    "test_labels: [n_rows, n_labels]",
    "return: []",
)
def loss(weights: ndarray, test_features: ndarray, test_labels: ndarray) -> ndarray:
    prediction: ndarray = check_shape(test_features @ weights, "[n_rows, n_labels]")
    error: ndarray = check_shape(prediction - test_labels, "[n_rows, n_labels]")
    square_error = check_shape(error ** 2, "[n_rows, n_labels]")
    mean_square_error = check_shape(np.mean(square_error, axis=-1), "[n_rows]")
    root_mean_square_error = check_shape(np.sqrt(mean_square_error), "[n_rows]")
    loss: ndarray = np.mean(root_mean_square_error)
    return loss

Checking shapes without a decorator

While the check_shapes() decorator is the recommend way to use this library, it is possible to use it without the decorator. In fact the decorator is just a wrapper around the class ShapeChecker, which can be used to check shapes directly:

def linear_model(features: ndarray, weights: ndarray) -> ndarray:
    checker = ShapeChecker()
    checker.check_shape(features, "[batch..., n_features]")
    checker.check_shape(weights, "[n_features]")
    prediction: ndarray = checker.check_shape(
        np.einsum("...i,i -> ...", features, weights), "[batch...]"
    )
    return prediction

You can use the function get_shape_checker() to get the ShapeChecker used by any immediately surrounding check_shapes() decorator.

Documenting shapes

The check_shapes() decorator rewrites the docstring (.__doc__) of the decorated function to add information about shapes, in a format compatible with Sphinx.

Only functions that already have a docstring will be updated. Functions that have no docstring at all will not have one added, this is so that we do not override a docstring that would have been inherited from a super class.

For example:

@check_shapes(
    "features: [batch..., n_features]",
    "weights: [n_features]",
    "return: [batch...]",
)
def linear_model(features: ndarray, weights: ndarray) -> ndarray:
    """
    Computes a prediction from a linear model.

    :param features: Data to make predictions from.
    :param weights: Model weights.
    :returns: Model predictions.
    """
    prediction: ndarray = np.einsum("...i,i -> ...", features, weights)
    return prediction

will have .__doc__:

"""
Computes a prediction from a linear model.

:param features:
    * **features** has shape [*batch*..., *n_features*].

    Data to make predictions from.
:param weights:
    * **weights** has shape [*n_features*].

    Model weights.
:returns:
    * **return** has shape [*batch*...].

    Model predictions.
"""

if you do not wish to have your docstrings rewritten, you can disable it with set_rewrite_docstrings():

set_rewrite_docstrings(None)

Supported types

This library has built-in support for checking the shapes of:

• Python built-in scalars: bool, int, float and str.

• Python built-in sequences: tuple and list.

• NumPy ndarrays.

• TensorFlow Tensors and Variables.

• TensorFlow Probability DeferredTensors, including TransformedVariable.

• JAX ndarrays.

• PyTorch Tensors.

Shapes of custom types

check_shapes uses the function get_shape() to extract the shape of an object. get_shape() will try to read a property of the object with name shape, which should have type Optional[Tuple[Optional[int], ...]]. If you have a custom type that does not have such a property you can use register_get_shape() to extend get_shape() to extract the shape of your type:

class LinearModel:
    @check_shapes(
        "weights: [n_features]",
    )
    def __init__(self, weights: ndarray) -> None:
        self._weights = weights

    @check_shapes(
        "self: [n_features]",
        "features: [batch..., n_features]",
        "return: [batch...]",
    )
    def predict(self, features: ndarray) -> ndarray:
        prediction: ndarray = np.einsum("...i,i -> ...", features, self._weights)
        return prediction

@register_get_shape(LinearModel)
def get_linear_model_shape(model: LinearModel, context: ErrorContext) -> Shape:
    shape: Shape = model._weights.shape
    return shape

@check_shapes(
    "model: [n_features]",
    "test_features: [n_rows, n_features]",
    "test_labels: [n_rows]",
    "return: []",
)
def loss(model: LinearModel, test_features: ndarray, test_labels: ndarray) -> float:
    prediction = model.predict(test_features)
    return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

check_shapes in stack traces

If you use check_shapes consistently you will have a lot of functions wrapped in check_shapes(). This means that if you have an error many of the stack frames in the trace back would belong to check_shapes. For the sake of more readable error messages check_shapes has code to hide itself from trace backs. If you do not like this you can disable this behaviour with set_drop_frames():

set_drop_frames(True)