Repository: mrphys/tensorflow-mri (language: python)

Issue:
Add complex initializers and convolutions



Discussion / hints:

Fix patch:
diff --git a/.gitattributes b/.gitattributes
index 1bccc1fa..8f4d9650 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1,1 +1,2 @@
+*.dat filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
diff --git a/.zenodo.json b/.zenodo.json
index 664aed6a..8f9b0097 100644
--- a/.zenodo.json
+++ b/.zenodo.json
@@ -1,9 +1,9 @@
 {
     "title": "TensorFlow MRI",
 
-    "description": "TensorFlow MRI is a collection of TensorFlow add-ons for computational MRI.",
+    "description": "A library of TensorFlow operators for computational magnetic resonance imaging.",
 
-    "keywords": ["TensorFlow", "MRI", "Magnetic Resonance Imaging", "Machine Learning", "ML", "Python"],
+    "keywords": ["TensorFlow", "MRI", "Magnetic Resonance Imaging", "Machine Learning", "ML", "Artificial Intelligence", "AI", "Deep Learning", "Python", "C++"],
 
     "upload_type": "software",
     
diff --git a/Makefile b/Makefile
index cbbebffd..e49d756f 100644
--- a/Makefile
+++ b/Makefile
@@ -34,8 +34,8 @@ docs: $(TARGET)
 	rm -rf tools/docs/_build
 	rm -rf tools/docs/_templates
 	rm -rf tools/docs/api_docs
-	$(PYTHON) tools/docs/generate_templates.py
-	$(PYTHON) tools/docs/generate_documents.py
+	$(PYTHON) tools/docs/create_templates.py
+	$(PYTHON) tools/docs/create_documents.py
 	$(MAKE) -C tools/docs dirhtml PY_VERSION=$(PY_VERSION)
 
 test: $(wildcard tensorflow_mri/python/ops/*.py)
diff --git a/README.rst b/README.rst
index e62b24e4..3e2464c3 100644
--- a/README.rst
+++ b/README.rst
@@ -46,9 +46,15 @@ TensorFlow MRI contains operators for:
   (`tfmri.convex <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/convex>`_):
   convex functions (quadratic, L1, L2, Tikhonov, total variation, etc.) and
   optimizers (ADMM).
+* Keras initializers
+  (`tfmri.initializers <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/initializers>`_):
+  neural network initializers, including support for complex-valued weights.
+* I/O (`tfmri.io <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/io>`_):
+  additional I/O functions potentially useful when working with MRI data.
 * Keras layers
   (`tfmri.layers <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/layers>`_):
-  layers and building blocks for neural networks.
+  layers and building blocks for neural networks, including support for
+  complex-valued weights, inputs and outputs.
 * Linear algebra
   (`tfmri.linalg <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/linalg>`_):
   linear operators specialized for image processing and MRI.
@@ -76,7 +82,7 @@ TensorFlow MRI contains operators for:
 * Unconstrained optimization
   (`tfmri.optimize <https://mrphys.github.io/tensorflow-mri/api_docs/tfmri/optimize>`_):
   gradient descent, L-BFGS.
-* Supporting array manipulation and math tasks. 
+* And more, e.g., supporting array manipulation and math tasks. 
 
 .. end-intro
 
@@ -105,6 +111,7 @@ versions of TensorFlow and TensorFlow MRI according to the table below.
 ======================  ========================  ============
 TensorFlow MRI Version  TensorFlow Compatibility  Release Date
 ======================  ========================  ============
+v0.18.0                 v2.8.x                    May 6, 2022
 v0.17.0                 v2.8.x                    Apr 22, 2022
 v0.16.0                 v2.8.x                    Apr 13, 2022
 v0.15.0                 v2.8.x                    Apr 1, 2022
@@ -130,21 +137,31 @@ v0.4.0                  v2.6.x                    Aug 18, 2021
 Documentation
 -------------
 
-Visit the `docs <https://mrphys.github.io/tensorflow-mri/>`_ for the API
-reference and examples of usage. 
+Visit the `docs <https://mrphys.github.io/tensorflow-mri/>`_ for guides,
+tutorials and the API reference. 
 
-Contributions
--------------
+Issues
+------
 
 If you use this package and something does not work as you expected, please
 `file an issue <https://github.com/mrphys/tensorflow-mri/issues/new>`_
-describing your problem. We will do our best to help.
+describing your problem. We're here to help!
+
+Credits
+-------
+
+If you like this software, star the repository! |stars|
 
-Contributions are very welcome. Please create a pull request if you would like
-to make a contribution.
+.. |stars| image:: https://img.shields.io/github/stars/mrphys/tensorflow-mri?style=social
+    :target: https://github.com/mrphys/tensorflow-mri/stargazers
 
-Citation
---------
+If you find this software useful in your research, you can cite TensorFlow MRI
+using its `Zenodo record <https://doi.org/10.5281/zenodo.5151590>`_.
+
+In the above link, scroll down to the "Export" section and select your favorite
+export format to get an up-to-date citation.
+
+Contributions
+-------------
 
-If you find this software useful in your work, please
-`cite us <https://doi.org/10.5281/zenodo.5151590>`_.
+Contributions of any kind are welcome! Open an issue or pull request to begin.
diff --git a/RELEASE.rst b/RELEASE.rst
index e6b5ba76..d1df5f4d 100644
--- a/RELEASE.rst
+++ b/RELEASE.rst
@@ -1,18 +1,40 @@
-Release 0.17.0
+Release 0.18.0
 ==============
 
-Breaking Changes
-----------------
+Major Features and Improvements
+-------------------------------
 
-* ``tfmri.signal``:
+* ``tfmri.convex``:
 
-    * Renamed keyword argument ``traj`` of ``filter_kspace`` and ``crop_kspace``
-      to ``trajectory``.
+  * All references to ``ndim`` in ``ConvexFunction`` and its subclasses have
+    been deprecated in favor of ``domain_dimension`` or ``domain_shape``. This
+    applies to constructor arguments, properties and methods.
+  * Refactored static and dynamic shape properties and methods to single-source
+    shape information on the ``_shape`` and ``_shape_tensor`` methods.
+    ``domain_dimension``, ``domain_dimension_tensor``, ``batch_shape`` and
+    ``batch_shape_tensor`` now just call ``shape`` and ``shape_tensor`` and
+    should not be overriden.
 
+* ``tfmri.initializers``:
+
+  * New module with initializers ``VarianceScaling``, ``GlorotNormal``,
+    ``GlorotUniform``, ``HeNormal``, ``HeUniform``, ``LecunNormal`` and
+    ``LecunUniform``. All initializers are drop-in replacements for their
+    Keras counterparts and support complex values.
+
+* ``tfmri.io``:
+
+  * New function ``parse_twix`` to parse TWIX RAID files (Siemens raw data).
+
+* ``tfmri.layers``:
+
+  * Added new layers ``Conv1D``, ``Conv2D`` and ``Conv3D``. All layers
+    are drop-in replacements for their Keras counterparts and support
+    complex values.
 
-Major Features and Improvements
--------------------------------
 
-* ``tfmri.plot``:
+Bug Fixes and Other Changes
+---------------------------
 
-    * New utility ``tiled_image``.
+* Using a new API import system, which should address some issues with
+  the autocomplete features of some editors.
diff --git a/tensorflow_mri/__about__.py b/tensorflow_mri/__about__.py
index 1c7088d0..91d3884f 100644
--- a/tensorflow_mri/__about__.py
+++ b/tensorflow_mri/__about__.py
@@ -29,7 +29,7 @@
 __summary__ = "A collection of TensorFlow add-ons for computational MRI."
 __uri__ = "https://github.com/mrphys/tensorflow-mri"
 
-__version__ = "0.17.0"
+__version__ = "0.18.0"
 
 __author__ = "Javier Montalt Tordera"
 __email__ = "javier.montalt@outlook.com"
diff --git a/tensorflow_mri/__init__.py b/tensorflow_mri/__init__.py
index 366993fc..97b2c61b 100644
--- a/tensorflow_mri/__init__.py
+++ b/tensorflow_mri/__init__.py
@@ -1,18 +1,8 @@
-# Copyright 2021 University College London. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"TensorFlow MRI."
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""TensorFlow MRI."""
+import os as _os
+import sys as _sys
 
 from tensorflow_mri.__about__ import *
 
@@ -32,10 +22,33 @@
 
 from tensorflow_mri import python
 
-# Import public API.
-from tensorflow_mri.python.util import api_util
+# Import submodules.
+from tensorflow_mri._api import callbacks
+from tensorflow_mri._api import coils
+from tensorflow_mri._api import convex
+from tensorflow_mri._api import image
+from tensorflow_mri._api import initializers
+from tensorflow_mri._api import io
+from tensorflow_mri._api import math
+from tensorflow_mri._api import layers
+from tensorflow_mri._api import linalg
+from tensorflow_mri._api import losses
+from tensorflow_mri._api import metrics
+from tensorflow_mri._api import optimize
+from tensorflow_mri._api import plot
+from tensorflow_mri._api import recon
+from tensorflow_mri._api import sampling
+from tensorflow_mri._api import signal
+from tensorflow_mri._api import summary
 
-for namespace in api_util.get_namespaces():
-  globals()[namespace] = api_util.import_namespace(namespace)
+# Make sure directory containing top level submodules is in
+# the __path__ so that "from tensorflow_mri.foo import bar" works.
+# We're using callbacks, but there's nothing special about that.
+_API_MODULE = _sys.modules[__name__].callbacks
+_tfmri_api_dir = _os.path.dirname(_os.path.dirname(_API_MODULE.__file__))
+_current_module = _sys.modules[__name__]
 
-del api_util
+if not hasattr(_current_module, '__path__'):
+  __path__ = [_tfmri_api_dir]
+elif _tfmri_api_dir not in __path__:
+  __path__.append(_tfmri_api_dir)
diff --git a/tensorflow_mri/_api/__init__.py b/tensorflow_mri/_api/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tensorflow_mri/_api/callbacks/__init__.py b/tensorflow_mri/_api/callbacks/__init__.py
new file mode 100644
index 00000000..e466fbae
--- /dev/null
+++ b/tensorflow_mri/_api/callbacks/__init__.py
@@ -0,0 +1,5 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Keras callbacks."""
+
+from tensorflow_mri.python.callbacks.tensorboard_callbacks import TensorBoardImages as TensorBoardImages
diff --git a/tensorflow_mri/_api/coils/__init__.py b/tensorflow_mri/_api/coils/__init__.py
new file mode 100644
index 00000000..06f5eacb
--- /dev/null
+++ b/tensorflow_mri/_api/coils/__init__.py
@@ -0,0 +1,8 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Parallel imaging operations."""
+
+from tensorflow_mri.python.ops.coil_ops import estimate_coil_sensitivities as estimate_sensitivities
+from tensorflow_mri.python.ops.coil_ops import combine_coils as combine_coils
+from tensorflow_mri.python.ops.coil_ops import compress_coils as compress_coils
+from tensorflow_mri.python.ops.coil_ops import CoilCompressorSVD as CoilCompressorSVD
diff --git a/tensorflow_mri/_api/convex/__init__.py b/tensorflow_mri/_api/convex/__init__.py
new file mode 100644
index 00000000..b6a1b1bb
--- /dev/null
+++ b/tensorflow_mri/_api/convex/__init__.py
@@ -0,0 +1,19 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Convex optimization operations."""
+
+from tensorflow_mri.python.ops.convex_ops import ConvexFunction as ConvexFunction
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionAffineMappingComposition as ConvexFunctionAffineMappingComposition
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionLinearOperatorComposition as ConvexFunctionLinearOperatorComposition
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionIndicatorBall as ConvexFunctionIndicatorBall
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionIndicatorL1Ball as ConvexFunctionIndicatorL1Ball
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionIndicatorL2Ball as ConvexFunctionIndicatorL2Ball
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionNorm as ConvexFunctionNorm
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionL1Norm as ConvexFunctionL1Norm
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionL2Norm as ConvexFunctionL2Norm
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionL2NormSquared as ConvexFunctionL2NormSquared
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionTikhonov as ConvexFunctionTikhonov
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionTotalVariation as ConvexFunctionTotalVariation
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionQuadratic as ConvexFunctionQuadratic
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionLeastSquares as ConvexFunctionLeastSquares
+from tensorflow_mri.python.ops.optimizer_ops import admm_minimize as admm_minimize
diff --git a/tensorflow_mri/_api/image/__init__.py b/tensorflow_mri/_api/image/__init__.py
new file mode 100644
index 00000000..893c4fcf
--- /dev/null
+++ b/tensorflow_mri/_api/image/__init__.py
@@ -0,0 +1,21 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Image processing operations."""
+
+from tensorflow_mri.python.ops.image_ops import psnr as psnr
+from tensorflow_mri.python.ops.image_ops import psnr2d as psnr2d
+from tensorflow_mri.python.ops.image_ops import psnr3d as psnr3d
+from tensorflow_mri.python.ops.image_ops import ssim as ssim
+from tensorflow_mri.python.ops.image_ops import ssim2d as ssim2d
+from tensorflow_mri.python.ops.image_ops import ssim3d as ssim3d
+from tensorflow_mri.python.ops.image_ops import ssim_multiscale as ssim_multiscale
+from tensorflow_mri.python.ops.image_ops import ssim2d_multiscale as ssim2d_multiscale
+from tensorflow_mri.python.ops.image_ops import ssim3d_multiscale as ssim3d_multiscale
+from tensorflow_mri.python.ops.image_ops import image_gradients as image_gradients
+from tensorflow_mri.python.ops.image_ops import gradient_magnitude as gradient_magnitude
+from tensorflow_mri.python.ops.image_ops import gmsd as gmsd
+from tensorflow_mri.python.ops.image_ops import gmsd2d as gmsd2d
+from tensorflow_mri.python.ops.image_ops import gmsd3d as gmsd3d
+from tensorflow_mri.python.ops.image_ops import total_variation as total_variation
+from tensorflow_mri.python.ops.image_ops import extract_glimpses as extract_glimpses
+from tensorflow_mri.python.ops.image_ops import phantom as phantom
diff --git a/tensorflow_mri/_api/initializers/__init__.py b/tensorflow_mri/_api/initializers/__init__.py
new file mode 100644
index 00000000..a1513d99
--- /dev/null
+++ b/tensorflow_mri/_api/initializers/__init__.py
@@ -0,0 +1,11 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Keras initializers."""
+
+from tensorflow_mri.python.initializers.initializers import VarianceScaling as VarianceScaling
+from tensorflow_mri.python.initializers.initializers import GlorotNormal as GlorotNormal
+from tensorflow_mri.python.initializers.initializers import GlorotUniform as GlorotUniform
+from tensorflow_mri.python.initializers.initializers import HeNormal as HeNormal
+from tensorflow_mri.python.initializers.initializers import HeUniform as HeUniform
+from tensorflow_mri.python.initializers.initializers import LecunNormal as LecunNormal
+from tensorflow_mri.python.initializers.initializers import LecunUniform as LecunUniform
diff --git a/tensorflow_mri/_api/io/__init__.py b/tensorflow_mri/_api/io/__init__.py
new file mode 100644
index 00000000..62d7219a
--- /dev/null
+++ b/tensorflow_mri/_api/io/__init__.py
@@ -0,0 +1,6 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Input/output operations."""
+
+from tensorflow_mri.python.io.image_io import encode_gif as encode_gif
+from tensorflow_mri.python.io.twix_io import parse_twix as parse_twix
diff --git a/tensorflow_mri/_api/layers/__init__.py b/tensorflow_mri/_api/layers/__init__.py
new file mode 100644
index 00000000..1dd6d65a
--- /dev/null
+++ b/tensorflow_mri/_api/layers/__init__.py
@@ -0,0 +1,9 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Keras layers."""
+
+from tensorflow_mri.python.layers.convolutional import Conv1D as Conv1D
+from tensorflow_mri.python.layers.convolutional import Conv2D as Conv2D
+from tensorflow_mri.python.layers.convolutional import Conv3D as Conv3D
+from tensorflow_mri.python.layers.conv_blocks import ConvBlock as ConvBlock
+from tensorflow_mri.python.layers.conv_endec import UNet as UNet
diff --git a/tensorflow_mri/_api/linalg/__init__.py b/tensorflow_mri/_api/linalg/__init__.py
new file mode 100644
index 00000000..6a2ab181
--- /dev/null
+++ b/tensorflow_mri/_api/linalg/__init__.py
@@ -0,0 +1,14 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Linear algebra operations."""
+
+from tensorflow_mri.python.util.linalg_imaging import LinearOperator as LinearOperator
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorAdjoint as LinearOperatorAdjoint
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorComposition as LinearOperatorComposition
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorAddition as LinearOperatorAddition
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorScaledIdentity as LinearOperatorScaledIdentity
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorDiag as LinearOperatorDiag
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorGramMatrix as LinearOperatorGramMatrix
+from tensorflow_mri.python.util.linalg_imaging import LinearOperatorFiniteDifference as LinearOperatorFiniteDifference
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorMRI as LinearOperatorMRI
+from tensorflow_mri.python.ops.linalg_ops import conjugate_gradient as conjugate_gradient
diff --git a/tensorflow_mri/_api/losses/__init__.py b/tensorflow_mri/_api/losses/__init__.py
new file mode 100644
index 00000000..ba51f107
--- /dev/null
+++ b/tensorflow_mri/_api/losses/__init__.py
@@ -0,0 +1,15 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Keras losses."""
+
+from tensorflow_mri.python.losses.confusion_losses import ConfusionLoss as ConfusionLoss
+from tensorflow_mri.python.losses.confusion_losses import FocalTverskyLoss as FocalTverskyLoss
+from tensorflow_mri.python.losses.confusion_losses import TverskyLoss as TverskyLoss
+from tensorflow_mri.python.losses.confusion_losses import F1Loss as F1Loss
+from tensorflow_mri.python.losses.confusion_losses import F1Loss as DiceLoss
+from tensorflow_mri.python.losses.confusion_losses import IoULoss as IoULoss
+from tensorflow_mri.python.losses.confusion_losses import IoULoss as JaccardLoss
+from tensorflow_mri.python.losses.iqa_losses import StructuralSimilarityLoss as StructuralSimilarityLoss
+from tensorflow_mri.python.losses.iqa_losses import MultiscaleStructuralSimilarityLoss as MultiscaleStructuralSimilarityLoss
+from tensorflow_mri.python.losses.iqa_losses import ssim_loss as ssim_loss
+from tensorflow_mri.python.losses.iqa_losses import ssim_multiscale_loss as ssim_multiscale_loss
diff --git a/tensorflow_mri/_api/math/__init__.py b/tensorflow_mri/_api/math/__init__.py
new file mode 100644
index 00000000..6164bc96
--- /dev/null
+++ b/tensorflow_mri/_api/math/__init__.py
@@ -0,0 +1,16 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Math operations."""
+
+from tensorflow_mri.python.ops.math_ops import extract_from_complex as extract_from_complex
+from tensorflow_mri.python.ops.math_ops import make_val_and_grad_fn as make_val_and_grad_fn
+from tensorflow_mri.python.ops.math_ops import normalize_no_nan as normalize_no_nan
+from tensorflow_mri.python.ops.math_ops import block_soft_threshold as block_soft_threshold
+from tensorflow_mri.python.ops.math_ops import shrinkage as shrinkage
+from tensorflow_mri.python.ops.math_ops import soft_threshold as soft_threshold
+from tensorflow_mri.python.ops.math_ops import indicator_box as indicator_box
+from tensorflow_mri.python.ops.math_ops import indicator_simplex as indicator_simplex
+from tensorflow_mri.python.ops.math_ops import indicator_ball as indicator_ball
+from tensorflow_mri.python.ops.math_ops import project_onto_box as project_onto_box
+from tensorflow_mri.python.ops.math_ops import project_onto_simplex as project_onto_simplex
+from tensorflow_mri.python.ops.math_ops import project_onto_ball as project_onto_ball
diff --git a/tensorflow_mri/_api/metrics/__init__.py b/tensorflow_mri/_api/metrics/__init__.py
new file mode 100644
index 00000000..3a81ef25
--- /dev/null
+++ b/tensorflow_mri/_api/metrics/__init__.py
@@ -0,0 +1,24 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Keras metrics."""
+
+from tensorflow_mri.python.metrics.confusion_metrics import ConfusionMetric as ConfusionMetric
+from tensorflow_mri.python.metrics.confusion_metrics import Accuracy as Accuracy
+from tensorflow_mri.python.metrics.confusion_metrics import TruePositiveRate as TruePositiveRate
+from tensorflow_mri.python.metrics.confusion_metrics import TruePositiveRate as Recall
+from tensorflow_mri.python.metrics.confusion_metrics import TruePositiveRate as Sensitivity
+from tensorflow_mri.python.metrics.confusion_metrics import TrueNegativeRate as TrueNegativeRate
+from tensorflow_mri.python.metrics.confusion_metrics import TrueNegativeRate as Specificity
+from tensorflow_mri.python.metrics.confusion_metrics import TrueNegativeRate as Selectivity
+from tensorflow_mri.python.metrics.confusion_metrics import PositivePredictiveValue as PositivePredictiveValue
+from tensorflow_mri.python.metrics.confusion_metrics import PositivePredictiveValue as Precision
+from tensorflow_mri.python.metrics.confusion_metrics import NegativePredictiveValue as NegativePredictiveValue
+from tensorflow_mri.python.metrics.confusion_metrics import TverskyIndex as TverskyIndex
+from tensorflow_mri.python.metrics.confusion_metrics import FBetaScore as FBetaScore
+from tensorflow_mri.python.metrics.confusion_metrics import F1Score as F1Score
+from tensorflow_mri.python.metrics.confusion_metrics import F1Score as DiceIndex
+from tensorflow_mri.python.metrics.confusion_metrics import IoU as IoU
+from tensorflow_mri.python.metrics.confusion_metrics import IoU as JaccardIndex
+from tensorflow_mri.python.metrics.iqa_metrics import PeakSignalToNoiseRatio as PeakSignalToNoiseRatio
+from tensorflow_mri.python.metrics.iqa_metrics import StructuralSimilarity as StructuralSimilarity
+from tensorflow_mri.python.metrics.iqa_metrics import MultiscaleStructuralSimilarity as MultiscaleStructuralSimilarity
diff --git a/tensorflow_mri/_api/optimize/__init__.py b/tensorflow_mri/_api/optimize/__init__.py
new file mode 100644
index 00000000..cdc4e3dc
--- /dev/null
+++ b/tensorflow_mri/_api/optimize/__init__.py
@@ -0,0 +1,6 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Optimization operations."""
+
+from tensorflow_mri.python.ops.optimizer_ops import gradient_descent as gradient_descent
+from tensorflow_mri.python.ops.optimizer_ops import lbfgs_minimize as lbfgs_minimize
diff --git a/tensorflow_mri/_api/plot/__init__.py b/tensorflow_mri/_api/plot/__init__.py
new file mode 100644
index 00000000..a3f95e60
--- /dev/null
+++ b/tensorflow_mri/_api/plot/__init__.py
@@ -0,0 +1,9 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Plotting utilities."""
+
+from tensorflow_mri.python.util.plot_util import plot_image_sequence as image_sequence
+from tensorflow_mri.python.util.plot_util import plot_tiled_image_sequence as tiled_image_sequence
+from tensorflow_mri.python.util.plot_util import plot_tiled_image as tiled_image
+from tensorflow_mri.python.util.plot_util import show as show
+from tensorflow_mri.python.util.plot_util import close as close
diff --git a/tensorflow_mri/_api/recon/__init__.py b/tensorflow_mri/_api/recon/__init__.py
new file mode 100644
index 00000000..8ce8d793
--- /dev/null
+++ b/tensorflow_mri/_api/recon/__init__.py
@@ -0,0 +1,9 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Image reconstruction."""
+
+from tensorflow_mri.python.ops.recon_ops import reconstruct_adj as adj
+from tensorflow_mri.python.ops.recon_ops import reconstruct_lstsq as lstsq
+from tensorflow_mri.python.ops.recon_ops import reconstruct_sense as sense
+from tensorflow_mri.python.ops.recon_ops import reconstruct_grappa as grappa
+from tensorflow_mri.python.ops.recon_ops import reconstruct_pf as pf
diff --git a/tensorflow_mri/_api/sampling/__init__.py b/tensorflow_mri/_api/sampling/__init__.py
new file mode 100644
index 00000000..b5f337c8
--- /dev/null
+++ b/tensorflow_mri/_api/sampling/__init__.py
@@ -0,0 +1,16 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""k-space sampling operations."""
+
+from tensorflow_mri.python.ops.traj_ops import density_grid as density_grid
+from tensorflow_mri.python.ops.traj_ops import random_sampling_mask as random_mask
+from tensorflow_mri.python.ops.traj_ops import radial_trajectory as radial_trajectory
+from tensorflow_mri.python.ops.traj_ops import spiral_trajectory as spiral_trajectory
+from tensorflow_mri.python.ops.traj_ops import radial_density as radial_density
+from tensorflow_mri.python.ops.traj_ops import estimate_radial_density as estimate_radial_density
+from tensorflow_mri.python.ops.traj_ops import radial_waveform as radial_waveform
+from tensorflow_mri.python.ops.traj_ops import estimate_density as estimate_density
+from tensorflow_mri.python.ops.traj_ops import flatten_trajectory as flatten_trajectory
+from tensorflow_mri.python.ops.traj_ops import flatten_density as flatten_density
+from tensorflow_mri.python.ops.traj_ops import expand_trajectory as expand_trajectory
+from tensorflow_mri.python.ops.traj_ops import expand_density as expand_density
diff --git a/tensorflow_mri/_api/signal/__init__.py b/tensorflow_mri/_api/signal/__init__.py
new file mode 100644
index 00000000..b27e2970
--- /dev/null
+++ b/tensorflow_mri/_api/signal/__init__.py
@@ -0,0 +1,12 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Signal processing operations."""
+
+from tensorflow_mri.python.ops.signal_ops import hann as hann
+from tensorflow_mri.python.ops.signal_ops import hamming as hamming
+from tensorflow_mri.python.ops.signal_ops import atanfilt as atanfilt
+from tensorflow_mri.python.ops.signal_ops import filter_kspace as filter_kspace
+from tensorflow_mri.python.ops.signal_ops import crop_kspace as crop_kspace
+from tensorflow_mri.python.ops.fft_ops import fftn as fft
+from tensorflow_mri.python.ops.fft_ops import ifftn as ifft
+from tensorflow_nufft.python.ops.nufft_ops import nufft as nufft
diff --git a/tensorflow_mri/_api/summary/__init__.py b/tensorflow_mri/_api/summary/__init__.py
new file mode 100644
index 00000000..d26dac6a
--- /dev/null
+++ b/tensorflow_mri/_api/summary/__init__.py
@@ -0,0 +1,5 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Tensorboard summaries."""
+
+from tensorflow_mri.python.summary.image_summary import gif as gif
diff --git a/tensorflow_mri/python/__init__.py b/tensorflow_mri/python/__init__.py
index b92a9dee..bfa70afd 100644
--- a/tensorflow_mri/python/__init__.py
+++ b/tensorflow_mri/python/__init__.py
@@ -15,6 +15,7 @@
 "TFMRI Python code."
 
 from tensorflow_mri.python import callbacks
+from tensorflow_mri.python import initializers
 from tensorflow_mri.python import io
 from tensorflow_mri.python import layers
 from tensorflow_mri.python import losses
diff --git a/tensorflow_mri/python/initializers/__init__.py b/tensorflow_mri/python/initializers/__init__.py
new file mode 100644
index 00000000..33ca9575
--- /dev/null
+++ b/tensorflow_mri/python/initializers/__init__.py
@@ -0,0 +1,35 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Keras initializers."""
+
+from tensorflow_mri.python.initializers import initializers
+
+
+TFMRI_INITIALIZERS = {
+    'VarianceScaling': initializers.VarianceScaling,
+    'GlorotNormal': initializers.GlorotNormal,
+    'GlorotUniform': initializers.GlorotUniform,
+    'HeNormal': initializers.HeNormal,
+    'HeUniform': initializers.HeUniform,
+    'LecunNormal': initializers.LecunNormal,
+    'LecunUniform': initializers.LecunUniform,
+    'variance_scaling': initializers.VarianceScaling,
+    'glorot_normal': initializers.GlorotNormal,
+    'glorot_uniform': initializers.GlorotUniform,
+    'he_normal': initializers.HeNormal,
+    'he_uniform': initializers.HeUniform,
+    'lecun_normal': initializers.LecunNormal,
+    'lecun_uniform': initializers.LecunUniform,
+}
diff --git a/tensorflow_mri/python/initializers/initializers.py b/tensorflow_mri/python/initializers/initializers.py
new file mode 100644
index 00000000..5d6c09c7
--- /dev/null
+++ b/tensorflow_mri/python/initializers/initializers.py
@@ -0,0 +1,289 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Keras initializers."""
+
+import math
+import string
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow_mri.python.util import api_util
+
+
+_PARTITION_SHAPE = 'partition_shape'
+_PARTITION_OFFSET = 'partition_offset'
+_ALLOWED_INITIALIZER_KWARGS = [_PARTITION_SHAPE, _PARTITION_OFFSET]
+
+
+EXTENSION_NOTE = string.Template("""
+
+  .. note::
+    This initializer can be used as a drop-in replacement for
+    `tf.keras.initializers.${name}`_. However, this one also supports
+    initialization of complex-valued weights. Simply pass `dtype='complex64'`
+    or `dtype='complex128'` to its `__call__` method.
+
+  .. _tf.keras.initializers.${name}: https://www.tensorflow.org/api_docs/python/tf/keras/initializers/${name}
+
+""")
+
+
+def complex_variance_scaling(base):
+  """Adds complex-valued support to Keras variance scaling initializers.
+
+  Args:
+    base: The base class to be extended. Must be a subclass of
+      `tf.keras.initializers.VarianceScaling`.
+
+  Returns:
+    A subclass of `base` that supports complex-valued initialization.
+
+  Raises:
+    ValueError: if `base` is not a subclass of
+    `tf.keras.initializers.VarianceScaling`.
+  """
+  if not issubclass(base, tf.keras.initializers.VarianceScaling):
+    raise ValueError(
+        f'Expected base class to be a subclass of '
+        f'`tf.keras.initializers.VarianceScaling`, but got {base}.')
+
+  # We override the initializer's __call__ method.
+  def __call__(self, shape, dtype=None, **kwargs):  # pylint: disable=invalid-name
+    """Returns a tensor object initialized as specified by the initializer.
+
+    Args:
+      shape: Shape of the tensor.
+      dtype: Optional dtype of the tensor. Only floating point and complex types
+        are supported. If not specified, `tf.keras.backend.floatx()` is used,
+        which defaults to `float32` unless you configured it otherwise (via
+        `tf.keras.backend.set_floatx(float_dtype)`)
+      **kwargs: Additional keyword arguments.
+
+    Returns:
+      A tensor object initialized as specified by the initializer.
+    """
+    # pylint: disable=protected-access,no-else-return
+    _validate_kwargs(self.__class__.__name__, kwargs)
+    dtype = _assert_float_or_complex_dtype(dtype)
+    scale = self.scale
+    fan_in, fan_out = _compute_fans(shape)
+    if _PARTITION_SHAPE in kwargs:  # pylint: disable=consider-using-get
+      shape = kwargs[_PARTITION_SHAPE]
+    # Compute required variance (in `scale`).
+    if self.mode == 'fan_in':
+      scale /= max(1., fan_in)
+    elif self.mode == 'fan_out':
+      scale /= max(1., fan_out)
+    else:
+      scale /= max(1., (fan_in + fan_out) / 2.)
+    if self.distribution == 'truncated_normal':
+      if dtype.is_complex:
+        # constant is stddev of complex standard normal truncated to 2
+        stddev = math.sqrt(scale) / .95311164380491208
+        return stddev * _complex_truncated_normal(
+            self._random_generator, shape, 2.0, dtype)
+      else:
+        # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+        stddev = math.sqrt(scale) / .87962566103423978
+        return self._random_generator.truncated_normal(
+            shape, 0.0, stddev, dtype)
+    elif self.distribution == 'untruncated_normal':
+      if dtype.is_complex:
+        stddev = math.sqrt(scale)
+        return stddev * _complex_normal(self._random_generator, shape, dtype)
+      else:
+        stddev = math.sqrt(scale)
+        return self._random_generator.random_normal(shape, 0.0, stddev, dtype)
+    else:
+      if dtype.is_complex:
+        stddev = math.sqrt(scale)
+        return stddev * _complex_uniform(self._random_generator, shape, dtype)
+      else:
+        limit = math.sqrt(3.0 * scale)
+        return self._random_generator.random_uniform(
+            shape, -limit, limit, dtype)
+
+  # Dynamically create a subclass of `base` with the same name as `base` and
+  # with the overriden `__call__` method.
+  subclass = type(base.__name__, (base,), {'__call__': __call__})
+
+  # Copy docs from the base class, adding the extra note.
+  docstring = base.__doc__
+  doclines = docstring.split('\n')
+  doclines[1:1] = EXTENSION_NOTE.substitute(name=base.__name__).splitlines()
+  subclass.__doc__ = '\n'.join(doclines)
+
+  return subclass
+
+
+# Define the variance scaling initializers. We use a composition of three
+# decorators:
+#  1. `complex_variance_scaling`: Adds complex-valued support to a Keras
+#     variance scaling initializer.
+#  2. `register_keras_serializable`: Registers the new initializer with the
+#     Keras serialization framework.
+#  3. `export`: Exports the new initializer to the TFMRI API.
+VarianceScaling = api_util.export("initializers.VarianceScaling")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.VarianceScaling)))
+
+
+GlorotNormal = api_util.export("initializers.GlorotNormal")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.GlorotNormal)))
+
+
+GlorotUniform = api_util.export("initializers.GlorotUniform")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.GlorotUniform)))
+
+
+HeNormal = api_util.export("initializers.HeNormal")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.HeNormal)))
+
+
+HeUniform = api_util.export("initializers.HeUniform")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.HeUniform)))
+
+
+LecunNormal = api_util.export("initializers.LecunNormal")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.LecunNormal)))
+
+
+LecunUniform = api_util.export("initializers.LecunUniform")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_variance_scaling(tf.keras.initializers.LecunUniform)))
+
+
+def _complex_uniform(rng, shape, dtype):
+  """Samples random values from a disk on the complex plane.
+
+  The sampled uniform distribution has zero mean and unit variance.
+
+  Args:
+    rng: A `keras.backend.RandomGenerator`.
+    shape: The shape of the output tensor.
+    dtype: The dtype of the output tensor. Must be complex.
+
+  Returns:
+    A tensor of shape `shape` and dtype `dtype`.
+  """
+  radius = tf.math.sqrt(rng.random_uniform(shape, 0.0, 2.0, dtype.real_dtype))
+  theta = rng.random_uniform(shape, 0.0, 2 * np.pi, dtype.real_dtype)
+  return tf.cast(radius, dtype) * tf.math.exp(tf.dtypes.complex(0.0, theta))
+
+
+def _complex_normal(rng, shape, dtype):
+  """Samples random values from normal distribution on the complex plane.
+
+  The sampled distribution has zero mean and unit variance.
+
+  Args:
+    rng: A `keras.backend.RandomGenerator`.
+    shape: The shape of the output tensor.
+    dtype: The dtype of the output tensor. Must be complex.
+
+  Returns:
+    A tensor of shape `shape` and dtype `dtype`.
+  """
+  sqrt2 = tf.math.sqrt(tf.constant(2.0, dtype=dtype.real_dtype))
+  real = rng.random_normal(shape, 0.0, 1.0, dtype=dtype.real_dtype) / sqrt2
+  imag = rng.random_normal(shape, 0.0, 1.0, dtype=dtype.real_dtype) / sqrt2
+  return tf.dtypes.complex(real, imag)
+
+
+def _complex_truncated_normal(rng, shape, upper, dtype):
+  """Samples random values from truncated normal on the complex plane.
+
+  The modulus is truncated to `upper`. The distribution has zero mean and unit
+  variance prior to the truncation.
+
+  Args:
+    rng: A `keras.backend.RandomGenerator`.
+    shape: The shape of the output tensor.
+    upper: The upper bound on the modulus (truncation).
+    dtype: The dtype of the output tensor. Must be complex.
+
+  Returns:
+    A tensor of shape `shape` and dtype `dtype`.
+  """
+  t = ((1 - tf.math.exp(tf.constant(-(upper ** 2), dtype.real_dtype))) *
+       rng.random_uniform(shape, dtype=dtype.real_dtype))
+  radius = tf.math.sqrt(-tf.math.log(1 - t))  # pylint: disable=invalid-unary-operand-type
+  theta = rng.random_uniform(shape, 0.0, 2 * np.pi, dtype.real_dtype)
+  return tf.cast(radius, dtype) * tf.math.exp(tf.dtypes.complex(0.0, theta))
+
+
+def _assert_float_or_complex_dtype(dtype):
+  """Validate and return floating or complex point type based on `dtype`.
+
+  `dtype` must be a floating point or complex type.
+
+  Args:
+    dtype: The data type to validate.
+
+  Returns:
+    Validated type.
+
+  Raises:
+    ValueError: if `dtype` is not a floating point type.
+  """
+  if dtype is None:
+    dtype = tf.keras.backend.floatx()
+  dtype = tf.as_dtype(dtype)
+  if not (dtype.is_floating or dtype.is_complex):
+    raise ValueError(f'Expected floating point type, got {dtype}.')
+  return dtype
+
+
+def _compute_fans(shape):
+  """Computes the number of input and output units for a weight shape.
+
+  Args:
+    shape: Integer shape tuple or TF tensor shape.
+
+  Returns:
+    A tuple of integer scalars (fan_in, fan_out).
+  """
+  if len(shape) < 1:  # Just to avoid errors for constants.
+    fan_in = fan_out = 1
+  elif len(shape) == 1:
+    fan_in = fan_out = shape[0]
+  elif len(shape) == 2:
+    fan_in = shape[0]
+    fan_out = shape[1]
+  else:
+    # Assuming convolution kernels (2D, 3D, or more).
+    # kernel shape: (..., input_depth, depth)
+    receptive_field_size = 1
+    for dim in shape[:-2]:
+      receptive_field_size *= dim
+    fan_in = shape[-2] * receptive_field_size
+    fan_out = shape[-1] * receptive_field_size
+  return int(fan_in), int(fan_out)
+
+
+def _validate_kwargs(cls_name, kwargs, support_partition=True):
+  for kwarg in kwargs:
+    if kwarg not in _ALLOWED_INITIALIZER_KWARGS:  # pylint: disable=no-else-raise
+      raise TypeError(f'Unknown keyword arguments: {kwarg}. Allowed keyword '
+                      f'arguments: {_ALLOWED_INITIALIZER_KWARGS}.')
+    elif not support_partition:
+      raise ValueError(f'{cls_name} initializer doesn\'t support '
+                       'partition-related arguments.')
diff --git a/tensorflow_mri/python/io/__init__.py b/tensorflow_mri/python/io/__init__.py
index fb4e0b15..44032b6c 100644
--- a/tensorflow_mri/python/io/__init__.py
+++ b/tensorflow_mri/python/io/__init__.py
@@ -15,3 +15,4 @@
 """Input/output."""
 
 from tensorflow_mri.python.io import image_io
+from tensorflow_mri.python.io import twix_io
diff --git a/tensorflow_mri/python/io/twix_io.py b/tensorflow_mri/python/io/twix_io.py
new file mode 100644
index 00000000..cb8c0799
--- /dev/null
+++ b/tensorflow_mri/python/io/twix_io.py
@@ -0,0 +1,1116 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""I/O operations with TWIX RAID files (Siemens raw data)."""
+
+import dataclasses
+import inspect
+import io
+import os
+import re
+import reprlib
+import struct
+import sys
+import typing
+import warnings
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow_mri.python.util import api_util
+
+NUL = b'\x00'
+
+
+@api_util.export("io.parse_twix")
+def parse_twix(contents):
+  """Parses the contents of a TWIX RAID file (Siemens raw data).
+
+  .. warning::
+    This function does not support graph execution.
+
+  Example:
+    >>> # Read bytes from file.
+    >>> contents = tf.io.read_file("/path/to/file.dat")
+    >>> # Parse the contents.
+    >>> twix = tfmri.io.parse_twix(contents)
+    >>> # Access the first measurement.
+    >>> meas = twix.measurements[0]
+    >>> # Get the protocol...
+    >>> protocol = meas.protocol
+    >>> # You can index the protocol to access any of the protocol buffers,
+    >>> # e.g., the measurement protocol.
+    >>> meas_prot = protocol['Meas']
+    >>> # Protocol buffers are nested structures accessible with "dot notation"
+    >>> # or "bracket notation". The following are equivalent:
+    >>> base_res = meas_prot.MEAS.sKSpace.lBaseResolution.value
+    >>> base_res = meas_prot['MEAS']['sKSpace']['lBaseResolution'].value
+    >>> # The measurement object also contains the scan data.
+    >>> scans = meas.scans
+    >>> # Each scan has a header and the list of channels.
+    >>> scan_header = scans[0].header
+    >>> channels = scans[0].channels
+    >>> # Each channel also has its own header as well as the raw measurement
+    >>> # data.
+    >>> channel_header = channels[0].header
+    >>> data = channels[0].data
+
+  Args:
+    contents: A scalar `tf.Tensor` of type `string`. The encoded contents of a
+      TWIX RAID file.
+
+  Returns:
+    A `TwixRaidFile` object containing the measurements, including protocol
+    buffers, scan headers and scan data.
+
+  Raises:
+    RuntimeError: If called with eager execution disabled.
+  """
+  if not tf.executing_eagerly():
+    raise RuntimeError('`parse_twix` must be called eagerly.')
+
+  # To scalar array.
+  if isinstance(contents, tf.Tensor):
+    contents = contents.numpy()
+
+  # Create a bytes stream.
+  stream = io.BytesIO(contents)
+
+  # Parse the object.
+  twix = TwixRaidFile.parse(stream)
+
+  # Close the stream. # TODO(jmontalt): This causes an error, so temporarily
+  # disabled.
+  # stream.close()
+
+  return twix
+
+
+@dataclasses.dataclass
+class ChannelHeader():
+  """Data structure for channel header."""
+  # `type` and `channel_length` are unpacked from ulTypeAndChannelLength.
+  # type_and_channel_length: int  # uint32
+  type: int  # uint8
+  channel_length:  int # uint24
+  meas_uid: int  # int32
+  scan_counter: int  # uint32
+  reserved1: int  # uint32
+  sequence_time: int  # uint32
+  unused2: int  # uint32
+  channel_id: int  # uint16
+  unused3: int  # uint16
+  crc: int  # uint32
+
+  @classmethod
+  def parse(cls, stream):
+    fmt = '<IiIIIIHHI'
+    values = struct.unpack(fmt, stream.read(struct.calcsize(fmt)))
+    type_and_channel_length = values[0]
+    type = type_and_channel_length & 0x000000ff  # pylint: disable=redefined-builtin
+    channel_length = type_and_channel_length >> 8
+    return cls(type, channel_length, *values[1:])
+
+CHANNEL_HEADER_SIZE = 32
+
+
+@dataclasses.dataclass
+class ChannelData():
+  """Data structure for channel data."""
+  header: ChannelHeader
+  data: tf.Tensor
+
+  @classmethod
+  def parse(cls, stream):
+    header = ChannelHeader.parse(stream)
+    return cls(header=header,
+               data=_read_tensor_data(
+                  stream, np.complex64,
+                  header.channel_length - CHANNEL_HEADER_SIZE))
+
+
+@dataclasses.dataclass
+class SlicePosVec():
+  """Data structure for slice position vector."""
+  sag: float  # float32
+  cor: float  # float32
+  tra: float  # float32
+
+  @classmethod
+  def parse(cls, stream):
+    return cls(*_read_float32(stream, 3))
+
+
+@dataclasses.dataclass
+class SliceData():
+  """Data structure for slice data."""
+  slice_pos_vec: SlicePosVec
+  quaternion: typing.Tuple[float, float, float, float]  # float32 * 4
+
+  @classmethod
+  def parse(cls, stream):
+    return cls(slice_pos_vec=SlicePosVec.parse(stream),
+               quaternion=struct.unpack('<4f', stream.read(16)))
+
+
+@dataclasses.dataclass
+class CutOffData():
+  """Data structure for cut-off data."""
+  pre: int    # uint16
+  post: int   # uint16
+
+  @classmethod
+  def parse(cls, stream):
+    return cls(*_read_uint16(stream, 2))
+
+
+@dataclasses.dataclass
+class LoopCounters():
+  """Data structure for loop counters."""
+  line: int  # uint16 (same for all counters)
+  acquisition: int
+  slice: int
+  partition: int
+  echo: int
+  phase: int
+  repetition: int
+  set: int
+  seg: int
+  ida: int
+  idb: int
+  idc: int
+  idd: int
+  ide: int
+
+  @classmethod
+  def parse(cls, stream):
+    return cls(*_read_uint16(stream, 14))
+
+
+@dataclasses.dataclass
+class EvalInfoMask():
+  """Data structure for evaluation info mask."""
+  # pylint: disable=invalid-name
+  ACQEND: bool  # 0: Last scan.
+  RTFEEDBACK: bool  # 1: Realtime feedback scan.
+  HPFEEDBACK: bool  # 2: High performance feedback scan.
+  ONLINE: bool  # 3: Processing should be done online.
+  OFFLINE: bool  # 4: Processing should be done offline.
+  SYNCDATA: bool  # 5: Readout contains synchronous data.
+  noname6: bool
+  noname7: bool
+  LASTSCANINCONCAT: bool  # 8: Last scan in concatenation.
+  noname9: bool
+  RAWDATACORRECTION: bool  # 10: Raw data correction scan.
+  LASTSCANINMEAS: bool  # 11: Last scan in measurement.
+  SCANSCALEFACTOR: bool  # 12: Scan-specific additional scale factor.
+  _2NDHADAMARPULSE: bool  # 13: 2nd RF excitation of HADAMAR
+  REFPHASESTABSCAN: bool  # 14: Reference phase stabilization scan.
+  PHASESTABSCAN: bool  # 15: Phase stabilization scan.
+  D3FFT: bool  # 16: Execute 3D FFT.
+  SIGNREV: bool  # 17: Sign reversal.
+  PHASEFFT: bool  # 18: Execute phase FFT.
+  SWAPPED: bool  # 19: Swapped phase/readout direction.
+  POSTSHAREDLINE: bool  # 20: Shared line.
+  PHASCOR: bool  # 21: Phase correction data.
+  PATREFSCAN: bool  # 22: Additional scan for PAT reference line/partition.
+  PATREFANDIMASCAN: bool  # 23: PAT reference and image scan.
+  REFLECT: bool   # 24: Reflect line.
+  NOISEADJSCAN: bool  # 25: Noise adjustment scan.
+  SHARENOW: bool  # 26: All lines acquired from the actual and previous.
+  LASTMEASUREDLINE: bool  # 27: Current line is the last measured line.
+  FIRSTSCANINSLICE: bool  # 28: First scan in slice (needed for time stamps).
+  LASTSCANINSLICE: bool  # 29: Last scan in slice (needed for time stamps).
+  TREFFECTIVEBEGIN: bool  # 30: Begin time stamp for TReff (triggered meas).
+  TREFFECTIVEEND: bool  # 31: End time stamp for TReff (triggered meas).
+  REF_POSITION: bool  # 32: Reference position for move during scan images.
+  SLC_AVERAGED: bool  # 33: Averaged slice for slice partial averaging scheme.
+  TAGFLAG1: bool  # 34: Adjust scan.
+  CT_NORMALIZE: bool  # 35: Scan for correction maps for TimCT-Prescan norm.
+  SCAN_FIRST: bool  # 36: First scan of a particular map.
+  SCAN_LAST: bool  # 37: Last scan of a particular map.
+  SLICE_ACCEL_REFSCAN: bool  # 38: Single-band refscan slice acc multiband acq.
+  SLICE_ACCEL_PHASCOR: bool  # 39: Additional phase corr for slice acc.
+  FIRST_SCAN_IN_BLADE: bool  # 40: First scan in blade.
+  LAST_SCAN_IN_BLADE: bool  # 41: Last scan in blade.
+  LAST_BLADE_IN_TR: bool  # 42: Last blade in TR interval.
+  noname43: bool
+  PACE: bool  # 44: PACE scan.
+  RETRO_LASTPHASE: bool  # 45: Last phase in heartbeat.
+  RETRO_ENDOFMEAS: bool  # 46: ADC at end of measurement.
+  RETRO_REPEATTHISHEARTBEAT: bool  # 47: Triggers repeat of current heartbeat.
+  RETRO_REPEATPREVHEARTBEAT: bool  # 48: Triggers repeat of previous heartbeat.
+  RETRO_ABORTSCANNOW: bool  # 49: Just abort everything.
+  RETRO_LASTHEARTBEAT: bool  # 50: This ADC is from last heart beat (dummy).
+  RETRO_DUMMYSCAN: bool  # 51: Just a dummy scan, throw away.
+  RETRO_ARRDETDISABLED: bool  # 52: Disables arrhythmia detection.
+  B1_CONTROLLOOP: bool  # 53: Readout to be used for B1 control loop.
+  SKIP_ONLINE_PHASCOR: bool  # 54: Skip online phase correction even if on.
+  SKIP_REGRIDDING: bool  # 55: Skip regridding even if on.
+  noname56: bool
+  noname57: bool
+  noname58: bool
+  noname59: bool
+  noname60: bool
+  WIP_1: bool  # 61: Scan for WIP application "type 1".
+  WIP_2: bool  # 62: Scan for WIP application "type 2".
+  WIP_3: bool  # 63: Scan for WIP application "type 3".
+
+  @classmethod
+  def parse(cls, stream):
+    mask = _read_uint64(stream)
+    num_flags = 64
+    bit_list = [bool((mask >> shift) & 1) for shift in range(num_flags)]
+    return cls(*bit_list)
+
+
+@dataclasses.dataclass
+class ScanHeader():
+  """Data structure for scan header."""
+  dma_length: int  # uint25
+  pack_bit: bool  # bool
+  flags: int  # uint32
+  meas_uid: int  # int32
+  scan_counter: int  # uint32
+  time_stamp: int  # uint32
+  pmu_time_stamp: int  # uint32
+  system_type: int  # uint16
+  ptab_pos_delay: int  # uint16
+  ptab_pos_x: int  # int32
+  ptab_pos_y: int  # int32
+  ptab_pos_z: int  # int32
+  reserved1: int  # uint32
+  eval_info_mask: EvalInfoMask  # uint64
+  samples_in_scan: int  # uint16
+  used_channels: int  # uint16
+  loop_counters: LoopCounters
+  cut_off: CutOffData
+  kspace_centre_column: int  # uint16
+  coil_select: int  # uint16
+  read_out_offcentre: float  # float32
+  time_since_last_rf: int  # uint32
+  kspace_centre_line_no: int  # uint16
+  kspace_centre_partition_no: int  # uint16
+  slice_data: SliceData
+  ice_program_para: typing.Tuple[int, ...]  # uint16 * 4
+  reserved_para: typing.Tuple[int, ...]  # uint16 * 4
+  application_counter: int  # uint16
+  application_mask: int  # uint16
+  crc: int  # uint32
+
+  @classmethod
+  def parse(cls, stream):
+    flags_and_dma_length = _read_uint32(stream)
+    return cls(dma_length=flags_and_dma_length % (2 ** 25),
+               pack_bit = bool((flags_and_dma_length >> 25) & 1),
+               flags = flags_and_dma_length >> 26,
+               meas_uid=_read_int32(stream),
+               scan_counter=_read_uint32(stream),
+               time_stamp=_read_uint32(stream),
+               pmu_time_stamp=_read_uint32(stream),
+               system_type=_read_uint16(stream),
+               ptab_pos_delay=_read_uint16(stream),
+               ptab_pos_x=_read_int32(stream),
+               ptab_pos_y=_read_int32(stream),
+               ptab_pos_z=_read_int32(stream),
+               reserved1=_read_uint32(stream),
+               eval_info_mask=EvalInfoMask.parse(stream),
+               samples_in_scan=_read_uint16(stream),
+               used_channels=_read_uint16(stream),
+               loop_counters=LoopCounters.parse(stream),
+               cut_off=CutOffData.parse(stream),
+               kspace_centre_column=_read_uint16(stream),
+               coil_select=_read_uint16(stream),
+               read_out_offcentre=_read_float32(stream),
+               time_since_last_rf=_read_uint32(stream),
+               kspace_centre_line_no=_read_uint16(stream),
+               kspace_centre_partition_no=_read_uint16(stream),
+               slice_data=SliceData.parse(stream),
+               ice_program_para=_read_uint16(stream, 24),
+               reserved_para=_read_uint16(stream, 4),
+               application_counter=_read_uint16(stream),
+               application_mask=_read_uint16(stream),
+               crc=_read_uint32(stream))
+
+SCAN_HEADER_SIZE = 192
+
+
+@dataclasses.dataclass
+class ScanData():
+  """Data structure for scan data.
+
+  Attributes:
+    header: A `ScanHeader`.
+    channels: A `tuple` of `ChannelData` instances.
+    sync_data: A `bytes` instance containing synchronization data. Only
+      set for SYNCDATA scans.
+  """
+  header: ScanHeader
+  channels: typing.Tuple[ChannelData, ...]
+  sync_data: bytes
+
+  @classmethod
+  def parse(cls, stream):
+    """Parses the scan data.
+
+    Args:
+      stream: A file-like object. The stream position must be at the beginning
+        of the object to be parsed.
+
+    Returns:
+      A `ScanData` instance.
+    """
+    header = ScanHeader.parse(stream)
+
+    channels = []
+    sync_data = None
+    if header.eval_info_mask.ACQEND:
+      # Nothing to read, move to end.
+      stream.seek(header.dma_length - SCAN_HEADER_SIZE, os.SEEK_CUR)
+    elif header.eval_info_mask.SYNCDATA:
+      # Read the sync data.
+      sync_data = stream.read(header.dma_length - SCAN_HEADER_SIZE)
+    else:
+      for _ in range(header.used_channels):
+        channels.append(ChannelData.parse(stream))
+
+    return cls(header=header, channels=channels, sync_data=sync_data)
+
+
+@dataclasses.dataclass
+class ParamBool():
+  """A parameter holding a boolean value.
+
+  Attributes:
+    value: The boolean value.
+    properties: A dictionary of properties.
+  """
+  value: bool
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses a bool parameter.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamBool` instance.
+
+    Raises:
+      ValueError: If the string is not a valid bool.
+    """
+    string = string.strip()
+    # Get properties.
+    string, properties = _get_properties(string)
+    # If string is empty, return None.
+    if not string:
+      return cls(value=None, properties=properties)
+    # Remove quotes if present.
+    if string[0] == '"' and string[-1] == '"':
+      string = string[1:-1]
+    # Now parse the string.
+    if string.lower() == 'true':
+      value = True
+    elif string.lower() == 'false':
+      value = False
+    else:
+      raise ValueError(f"Invalid boolean value: {string}")
+    return cls(value=value, properties=properties)
+
+
+@dataclasses.dataclass
+class ParamChoice():
+  """A parameter holding a choice value.
+
+  Attributes:
+    value: The choice value.
+    properties: A dictionary of properties.
+  """
+  value: typing.Any
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses a choice parameter.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamChoice` instance.
+    """
+    string = string.strip()
+    # Get properties.
+    string, properties = _get_properties(string)
+    # If string is empty, return None.
+    if not string:
+      return cls(value=None, properties=properties)
+    # Remove quotes if present.
+    if string[0] == '"' and string[-1] == '"':
+      string = string[1:-1]
+    return cls(value=string, properties=properties)
+
+
+@dataclasses.dataclass
+class ParamDouble():
+  """A parameter holding one or more real values.
+
+  Attributes:
+    value: The real value. May be a tuple.
+    properties: A dictionary of properties.
+  """
+  value: typing.Union[float, typing.Tuple[float, ...]]
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses a real parameter.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamDouble` instance.
+    """
+    string = string.strip()
+    # Get properties.
+    string, properties = _get_properties(string)
+    # If string is empty, return None.
+    if not string:
+      return cls(value=None, properties=properties)
+    value = string.split()
+    if len(value) > 1:
+      value = tuple(float(s.strip()) for s in value)
+    else:
+      value = float(value[0])
+    return cls(value=value, properties=properties)
+
+
+@dataclasses.dataclass
+class ParamLong():
+  """A parameter holding one or more integer values.
+
+  Attributes:
+    value: The integer value. May be a tuple.
+    properties: A dictionary of properties.
+  """
+  value: typing.Union[int, typing.Tuple[int, ...]]
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses an integer parameter.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamLong` instance.
+    """
+    string = string.strip()
+    # Get properties.
+    string, properties = _get_properties(string)
+    # If string is empty, return None.
+    if not string:
+      return cls(value=None, properties=properties)
+    # Support single inputs and arrays.
+    value = string.split()
+    if len(value) > 1:
+      value = tuple(int(s.strip()) for s in value)
+    else:
+      value = int(value[0])
+    return cls(value=value, properties=properties)
+
+
+@dataclasses.dataclass
+class ParamString():
+  """A parameter holding a string value.
+
+  Attributes:
+    value: The string value.
+    properties: A dictionary of properties.
+  """
+  value: str
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses a string parameter.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamString` instance.
+    """
+    string = string.strip()
+    # Get properties.
+    string, properties = _get_properties(string)
+    # If empty, we're done.
+    if not string:
+      return cls(value=string, properties=properties)
+    # Remove quotes if present.
+    if string[0] == '"' and string[-1] == '"':
+      string = string[1:-1]
+    return cls(value=string, properties=properties)
+
+
+@dataclasses.dataclass
+class ParamMap():
+  """A parameter map.
+
+  Entries in the `params` dictionary can be directly accessed via item
+  or attr getters, i.e., both `param_map['key']` and `param_map.key` are
+  supported and equivalent to `param_map.params['key']`.
+
+  Attributes:
+    params: A `dict` of parameters.
+    properties: A `dict` of properties.
+  """
+  params: typing.Mapping[str, typing.Any]
+  properties: typing.Mapping[str, typing.Any]
+
+  @classmethod
+  def parse(cls, string):
+    """Parses a parameter map.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ParamMap` instance.
+
+    Raises:
+      ValueError: If the string is invalid.
+    """
+    params = {}
+    string, properties = _get_properties(string)
+
+    while string.strip():
+      string, param = _get_next_param(string)
+      if param is None:
+        raise ValueError(
+            f"Could not parse next parameter in string:\n "
+            f"{_head_string(string)}...")
+
+      type_, name, value = param
+      if name is not None:
+        if type_ == 'ParamBool':
+          params[name] = ParamBool.parse(value)
+        elif type_ == 'ParamChoice':
+          params[name] = ParamChoice.parse(value)
+        elif type_ == 'ParamDouble':
+          params[name] = ParamDouble.parse(value)
+        elif type_ == 'ParamLong':
+          params[name] = ParamLong.parse(value)
+        elif type_ == 'ParamMap':
+          params[name] = ParamMap.parse(value)
+        elif type_ == 'ParamString':
+          params[name] = ParamString.parse(value)
+        else:
+          warnings.warn(f'Unknown param type {type_} will not be parsed.')
+          params[name] = value
+    return cls(params, properties)
+
+  def __getitem__(self, name):
+    return self.params[name]
+
+  def __getattr__(self, name):
+    if name not in self.params:
+      raise AttributeError(f'Parameter `{name}` does not exist.')
+    return self.params[name]
+
+
+class ProtocolBuffer(ParamMap):
+  """Data structure for a protocol buffer.
+
+  This is a special type of `ParamMap`.
+  """
+  @classmethod
+  def parse(cls, string):
+    """Parses a protocol buffer.
+
+    Args:
+      string: The `str` to parse.
+
+    Returns:
+      A `ProtocolBuffer` instance.
+
+    Raises:
+      ValueError: If the string is not a valid protocol buffer.
+    """
+    original = string  # For error messages.
+
+    # Get buffer properties.
+    string, properties = _get_properties(string)
+
+    # Get the top-level parameter map and begin recursive parsing.
+    string, (type_, name, value) = _get_next_param(string)
+    if type_ != 'ParamMap' or name != '':
+      raise ValueError(f'Invalid ProtocolBuffer string: {original}')
+    param_map = ParamMap.parse(value)
+
+    return cls(param_map.params, properties)
+
+
+@dataclasses.dataclass
+class MeasurementHeader():
+  """Data structure for a measurement header.
+
+  Attributes:
+    meas_id: The measurement ID.
+    file_id: The file ID.
+    offset: The start offset of the corresponding measurement, in bytes.
+    length: The length of the corresponding measurement, in bytes.
+    pat_name: The name of the patient.
+    prot_name: The name of the protocol.
+  """
+  meas_id: int  # uint32
+  file_id: int  # uint32
+  offset: int  # uint64
+  length: int  # uint64
+  pat_name: str  # string (64 bytes)
+  prot_name: str  # string (64 bytes)
+
+  @classmethod
+  def parse(cls, stream):
+    """Parses a `MeasurementHeader` from a stream.
+
+    Args:
+      stream: A file-like object. The stream position must be at the beginning
+        of the object to be parsed.
+
+    Returns:
+      A `MeasurementHeader` object.
+    """
+    meas_id = _read_uint32(stream)
+    file_id = _read_uint32(stream)
+    offset = _read_uint64(stream)
+    length = _read_uint64(stream)
+    pat_name = _read_string(stream, 64)
+    prot_name = _read_string(stream, 64)
+
+    return cls(meas_id=meas_id, file_id=file_id, offset=offset, length=length,
+               pat_name=pat_name, prot_name=prot_name)
+
+
+@dataclasses.dataclass
+class MeasurementData():
+  """Data structure for a measurement.
+
+  Attributes:
+    header: A `MeasurementHeader`.
+    protocol: A `dict` of `ProtocolBuffer` objects.
+    scans: A `tuple` of `ScanData` objects.
+  """
+  header: MeasurementHeader
+  protocol: typing.Mapping[str, ProtocolBuffer]
+  scans: typing.Tuple[ScanData, ...]
+
+  @classmethod
+  def parse(cls, stream, header):
+    """Parses a `MeasurementData` from a stream.
+
+    Args:
+      stream: A file-like object. The stream position must be at the beginning
+        of the object to be parsed.
+      header: A `MeasurementHeader` object.
+
+    Returns:
+      A `MeasurementData` object.
+    """
+    # Place stream position at the beginning of the measurement.
+    stream.seek(header.offset)
+
+    # Some header info.
+    prot_length = _read_uint32(stream)
+    num_buffers = _read_uint32(stream)
+
+    # Read protocol buffers.
+    protocol = {}
+    for _ in range(num_buffers):
+      # Buffer header.
+      buffer_name = _read_null_terminated_string(stream)
+      buffer_length = _read_uint32(stream)
+
+      # Read, decode and remove NUL terminator.
+      buffer_string = stream.read(buffer_length).decode('latin1')[:-1]
+      # Parse the protocol buffer string.
+      protocol[buffer_name] = _parse_protocol_buffer(buffer_string)
+
+    # Move to end of protocol.
+    stream.seek(header.offset + prot_length)
+
+    # Read scan data.
+    scans = []
+    while True:  # Continue until ACQEND.
+      scans.append(ScanData.parse(stream))
+
+      # If ACQEND is set, we're done.
+      if scans[-1].header.eval_info_mask.ACQEND:
+        break
+
+    return cls(header=header, protocol=protocol, scans=scans)
+
+
+MAX_RAID_FILE_MEASUREMENTS = 64
+RAID_FILE_HEADER_SIZE = 152
+
+
+@dataclasses.dataclass
+class TwixRaidFileHeader():
+  """Data structure for a TWIX raid file header.
+
+  Attributes:
+    meas_count: Number of measurements in the file.
+    meas_headers: A `tuple` of `MeasurementHeader` objects.
+  """
+  meas_count: int  # uint32
+  meas_headers: typing.Tuple[MeasurementHeader, ...]
+
+  @classmethod
+  def parse(cls, stream):
+    """Parses a `TwixRaidFileHeader` from a bytes stream.
+
+    Args:
+      stream: A file-like object. The stream position must be at the beginning
+        of the object to be parsed.
+
+    Returns:
+      A `TwixRaidFileHeader` instance.
+    """
+    _ = _read_uint32(stream)  # Ignored.
+    meas_count = _read_uint32(stream)
+
+    meas_headers = []
+    for _ in range(meas_count):
+      meas_headers.append(MeasurementHeader.parse(stream))
+
+    unused_entries = MAX_RAID_FILE_MEASUREMENTS - meas_count
+    stream.seek(unused_entries * RAID_FILE_HEADER_SIZE, io.SEEK_CUR)
+
+    return cls(meas_count=meas_count, meas_headers=meas_headers)
+
+
+@dataclasses.dataclass
+class TwixRaidFile():
+  """Data structure for a TWIX raid file.
+
+  Attributes:
+    measurements: A `tuple` of `MeasurementData` objects.
+  """
+  measurements: typing.Tuple[MeasurementData, ...]
+
+  @classmethod
+  def parse(cls, stream):
+    """Parses a `TwixRaidFile` from a bytes stream.
+
+    Args:
+      stream: A file-like object. The stream position must be at the beginning
+        of the object to be parsed.
+
+    Returns:
+      A `TwixRaidFile` instance.
+    """
+    header = TwixRaidFileHeader.parse(stream)
+
+    measurements = []
+    for entry in header.meas_headers:
+      stream.seek(entry.offset)
+      measurements.append(MeasurementData.parse(stream, entry))
+
+    return cls(measurements=measurements)
+
+
+def _parse_protocol_buffer(string):
+  """Parse a protocol buffer from text.
+
+  Args:
+    string: A `str`.
+
+  Returns:
+    A `ProtocolBuffer` instance.
+
+  Raises:
+    ValueError: If the string is not a valid protocol buffer.
+  """
+  original = string
+  # We can't parse ASCCONV for now, so remove them.
+  for m in RE_ASCCONV_RANGE.finditer(string):
+    string = string.replace(m.group(), '')
+
+  # String might be empty after removing ASCCONV. If so, just return an empty
+  # ProtocolBuffer.
+  string = string.strip()
+  if not string:
+    return ProtocolBuffer(params={}, properties={})
+
+  string, property_ = _get_next_property(string)
+  if property_ is None:
+    raise ValueError(
+        f"Failed while parsing protocol buffer:\n "
+        f"{_head_string(original)}")
+
+  name, value = property_
+  if name != 'XProtocol':
+    raise ValueError(
+        f"Failed while parsing protocol buffer:\n "
+        f"{_head_string(original)}")
+
+  return ProtocolBuffer.parse(value)
+
+
+def _read_int16(stream, count=None):
+  return _read_values(stream, 'h', count)
+
+def _read_uint16(stream, count=None):
+  return _read_values(stream, 'H', count)
+
+def _read_int32(stream, count=None):
+  return _read_values(stream, 'i', count)
+
+def _read_uint32(stream, count=None):
+  return _read_values(stream, 'I', count)
+
+def _read_int64(stream, count=None):
+  return _read_values(stream, 'q', count)
+
+def _read_uint64(stream, count=None):
+  return _read_values(stream, 'Q', count)
+
+def _read_float32(stream, count=None):
+  return _read_values(stream, 'f', count)
+
+def _read_float64(stream, count=None):
+  return _read_values(stream, 'd', count)
+
+def _read_values(stream, fmt, count):
+  fmt = f'<{count or 1}{fmt}'
+  values = struct.unpack(fmt, stream.read(struct.calcsize(fmt)))
+  if count is None:
+    values = values[0]
+  return values
+
+def _read_tensor_data(stream, dtype, length):
+  dtype = np.dtype(dtype)
+  # Read data from stream.
+  data = np.frombuffer(stream.getbuffer()[stream.tell():],
+                       dtype=dtype, count=length // dtype.itemsize)
+  # Move stream to end of data.
+  stream.seek(data.nbytes, io.SEEK_CUR)
+  return tf.convert_to_tensor(data)
+
+def _read_string(stream, length):
+  return stream.read(length).rstrip(NUL).decode('latin1')
+
+def _read_null_terminated_string(stream):
+  string = b''
+  while True:
+    char = stream.read(1)
+    if char == NUL:
+      break
+    string += char
+  return string.decode('latin1')
+
+# The following regex patterns are used to parse the protocol buffers.
+# Matches strings such as:
+#   <Visible>
+#   <ParamMap."">
+#   <ParamBool."TestParam">
+#   <ParamBool."Test;Param">
+#   <ParamBool."Test@Param">
+#   <ParamBool."Test Param">
+RE_TAG = re.compile(r'<(?P<type>\w+)(?:\."(?P<name>[\w@; ]*)")?>')
+# Matches only strings like:
+#   <Visible>
+RE_TAG_PROPERTY = re.compile(r'<(?P<name>\w+)>')
+# Matches only strings like:
+#   <ParamMap."">
+#   <ParamBool."TestParam">
+RE_TAG_PARAM = re.compile(r'<(?P<type>\w+)\."(?P<name>[\w@; ]*)">')
+# Matches an ASCCONV range (everything between ASCCONV BEGIN AND ASCCONV END).
+RE_ASCCONV_RANGE = re.compile(
+    r'### ASCCONV BEGIN(?P<content>(?:.|\n)*)### ASCCONV END ###')
+# Matches any whitespace character.
+RE_WHITESPACE = re.compile(r'\s')
+# Matches any non-whitespace character.
+RE_NON_WHITESPACE = re.compile(r'\S')
+# Matches a newline.
+RE_NEWLINE = re.compile(r'\n')
+
+
+def _get_next_element(string, re_tag):
+  """Returns the next element tagged with `re_tag`.
+
+  Args:
+    string: The string to parse.
+    re_tag: The regex pattern describing the tag.
+
+  Returns:
+    The string without the matched element and a tuple describing the matched
+    element.
+
+  Raises:
+    ValueError: If the string is not valid.
+  """
+  string = string.strip()
+  m = re_tag.match(string)
+  if not m:
+    return string, None
+  try:
+    type_ = m.group('type')
+  except IndexError:
+    # It's a property.
+    type_ = None
+  name = m.group('name')
+
+  # Remove tag from string.
+  string = string[m.end():]
+
+  # Get first non-whitespace character.
+  m = RE_NON_WHITESPACE.search(string)
+  if not m:
+    raise ValueError('Invalid string: {}'.format(string))
+  c = m.group()
+  string = string[m.start():]
+
+  # Now, depending on what character we found:
+  if name == "Comment":
+    # Continue until the end of the line.
+    m = RE_NEWLINE.search(string)
+    end = m.start() if m else len(string)
+    value = string[:end].strip()
+  elif c == '{':
+    # Find the matching closing brace.
+    level = 1
+    string = string[1:]
+    end = None
+    for end, c in enumerate(string):
+      if c == '{':
+        level += 1
+      elif c == '}':
+        level -= 1
+        if level == 0:
+          break
+    value = string[:end].strip()
+  elif c == '"':
+    # Find the matching closing quote.
+    string = string[1:]
+    end = None
+    for end, c in enumerate(string):
+      if c == '"':
+        break
+    value = string[:end].strip()
+  else:
+    # Continue until next whitespace or until the end of the string.
+    m = RE_WHITESPACE.search(string)
+    end = m.start() if m else len(string)
+    value = string[:end].strip()
+
+  if len(string) > 0:
+    string = string[end+1:]
+
+  if type_ is not None:
+    return string, (type_, name, value)
+  return string, (name, value)
+
+
+def _get_next_param(string):
+  """Returns the next parameter in string.
+
+  Args:
+    string: The string to parse.
+
+  Returns:
+    The string without the matched parameter and a tuple describing
+    the matched parameter (its type, name and value). If no parameter
+    is found, returns (string, None).
+  """
+  return _get_next_element(string, RE_TAG_PARAM)
+
+
+def _get_next_property(string):
+  """Returns the next property in string.
+
+  Args:
+    string: The string to parse.
+
+  Returns:
+    The string without the matched property and a tuple describing
+    the matched property (its name and value). If no property is found,
+    returns (string, None).
+  """
+  return _get_next_element(string, RE_TAG_PROPERTY)
+
+
+def _get_properties(string):
+  """Returns a list of properties found in string.
+
+  Args:
+    string: The string to parse.
+
+  Returns:
+    The string without the matched properties and a list of tuples
+    describing the matched properties (their name and value).
+  """
+  properties = {}
+  while True:
+    string, property_ = _get_next_property(string)
+    if property_ is None:
+      break
+    name, value = property_
+    if name == 'Precision':
+      value = int(value.strip())
+    else:
+      value = value.strip()
+    properties[name] = value
+  return string, properties
+
+
+def _head_string(string, n=10):
+  """Returns the first n lines of a string.
+
+  Args:
+    string: A `str`.
+    n: The number of lines to return.
+
+  Returns:
+    A `str` containing the first `n` lines of `string`.
+  """
+  return '\n'.join(string.splitlines()[:n])
+
+
+# Some of the objects in this module are deeply nested dataclasses which can
+# potentially return absurdly long string representations when printed.
+# To prevent this, we customise the string representations to limit the size
+# of the reprs for each field.
+REPR = reprlib.Repr()
+
+def __repr__(self):  # pylint: disable=invalid-name
+  """Returns a string representation of the object.
+
+  This repr is used for all dataclasses in this module. It is very similar to
+  the default repr, except that it limits the size of the reprs for the
+  field values.
+  """
+  fields = dataclasses.fields(self)
+  return (self.__class__.__qualname__ +
+          f"({', '.join(f'{f.name}={REPR.repr(getattr(self, f.name))}' for f in fields)})")  # pylint: disable=line-too-long
+
+# Set the compressed repr for all dataclasses in this module.
+for _, class_ in inspect.getmembers(sys.modules[__name__],
+                                    dataclasses.is_dataclass):
+  class_.__repr__ = __repr__
diff --git a/tensorflow_mri/python/layers/__init__.py b/tensorflow_mri/python/layers/__init__.py
index e4d8d0e5..feb6c488 100644
--- a/tensorflow_mri/python/layers/__init__.py
+++ b/tensorflow_mri/python/layers/__init__.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 """Keras layers."""
 
+from tensorflow_mri.python.layers import convolutional
 from tensorflow_mri.python.layers import conv_blocks
 from tensorflow_mri.python.layers import conv_endec
 from tensorflow_mri.python.layers import preproc_layers
diff --git a/tensorflow_mri/python/layers/convolutional.py b/tensorflow_mri/python/layers/convolutional.py
new file mode 100644
index 00000000..14830d2e
--- /dev/null
+++ b/tensorflow_mri/python/layers/convolutional.py
@@ -0,0 +1,133 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Convolutional layers."""
+
+import string
+
+import tensorflow as tf
+
+from tensorflow_mri.python.initializers import TFMRI_INITIALIZERS
+from tensorflow_mri.python.util import api_util
+
+
+EXTENSION_NOTE = string.Template("""
+
+  .. note::
+    This layer can be used as a drop-in replacement for
+    `tf.keras.layers.${name}`_. However, this one also supports complex-valued
+    convolutions. Simply pass `dtype='complex64'` or `dtype='complex128'` to
+    the layer constructor.
+
+  .. _tf.keras.layers.${name}: https://www.tensorflow.org/api_docs/python/tf/keras/layers/${name}
+
+""")
+
+
+def complex_conv(base):
+  """Adds complex-valued support to Keras conv layers.
+
+  We need two override two methods in the conv layer:
+
+  1. `__init__`: This allows us to replace some of the Keras initializers
+      that do not support complex dtypes by the corresponding TFMRI versions
+      with support initialization of complex-valued variables.
+  2. `convolution_op`: This implements the complex-valued convolution.
+
+  Args:
+    base: The base class to be extended. Must be one of
+      `tf.keras.layers.Conv1D`, `tf.keras.layers.Conv2D`, or
+      `tf.keras.layers.Conv3D`.
+
+  Returns:
+    A subclass of `base` that supports complex-valued convolutions.
+
+  Raises:
+    ValueError: If `base` is not one of the supported base classes.
+  """
+  if not issubclass(base, (tf.keras.layers.Conv1D,
+                           tf.keras.layers.Conv2D,
+                           tf.keras.layers.Conv3D)):
+    raise ValueError(
+        f'Expected base class to be a subclass of '
+        f'`tf.keras.layers.ConvND`, but got {base}.')
+
+  def __init__(self, *args, **kwargs):  # pylint: disable=invalid-name
+    # If the requested initializer is one of those provided by TFMRI, prefer
+    # the TFMRI version.
+    kernel_initializer = kwargs.get('kernel_initializer', 'glorot_uniform')
+    if (isinstance(kernel_initializer, str) and
+        kernel_initializer in TFMRI_INITIALIZERS):
+      kwargs['kernel_initializer'] = TFMRI_INITIALIZERS[kernel_initializer]()
+
+    bias_initializer = kwargs.get('bias_initializer', 'zeros')
+    if (isinstance(bias_initializer, str) and
+        bias_initializer in TFMRI_INITIALIZERS):
+      kwargs['bias_initializer'] = TFMRI_INITIALIZERS[bias_initializer]()
+
+    return base.__init__(self, *args, **kwargs)
+
+  def convolution_op(self, inputs, kernel):
+    # If dtype is not complex, fall back to original implementation.
+    if not tf.as_dtype(self.dtype).is_complex:
+      return base.convolution_op(self, inputs, kernel)
+
+    inputs_real = tf.math.real(inputs)
+    inputs_imag = tf.math.imag(inputs)
+    kernel_real = tf.math.real(kernel)
+    kernel_imag = tf.math.imag(kernel)
+
+    outputs_real = (
+        base.convolution_op(self, inputs_real, kernel_real) -
+        base.convolution_op(self, inputs_imag, kernel_imag))
+
+    outputs_imag = (
+        base.convolution_op(self, inputs_real, kernel_imag) +
+        base.convolution_op(self, inputs_imag, kernel_real))
+
+    return tf.dtypes.complex(outputs_real, outputs_imag)
+
+  # Dynamically create a subclass of `base` with the same name as `base` and
+  # with the overriden `convolution_op` method.
+  subclass = type(base.__name__, (base,), {'__init__': __init__,
+                                           'convolution_op': convolution_op})
+
+  # Copy docs from the base class, adding the extra note.
+  docstring = base.__doc__
+  doclines = docstring.split('\n')
+  doclines[1:1] = EXTENSION_NOTE.substitute(name=base.__name__).splitlines()
+  subclass.__doc__ = '\n'.join(doclines)
+
+  return subclass
+
+
+# Define the complex-valued conv layers. We use a composition of three
+# decorators:
+#  1. `complex_conv`: Adds complex-valued support to a Keras conv layer.
+#  2. `register_keras_serializable`: Registers the new layer with the Keras
+#     serialization framework.
+#  3. `export`: Exports the new layer to the TFMRI API.
+Conv1D = api_util.export("layers.Conv1D")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_conv(tf.keras.layers.Conv1D)))
+
+
+Conv2D = api_util.export("layers.Conv2D")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_conv(tf.keras.layers.Conv2D)))
+
+
+Conv3D = api_util.export("layers.Conv3D")(
+    tf.keras.utils.register_keras_serializable(package='MRI')(
+    complex_conv(tf.keras.layers.Conv3D)))
diff --git a/tensorflow_mri/python/ops/convex_ops.py b/tensorflow_mri/python/ops/convex_ops.py
index 7f4981eb..0083ab9f 100644
--- a/tensorflow_mri/python/ops/convex_ops.py
+++ b/tensorflow_mri/python/ops/convex_ops.py
@@ -21,6 +21,7 @@
 import tensorflow as tf
 
 from tensorflow_mri.python.ops import array_ops
+from tensorflow_mri.python.util import deprecation
 from tensorflow_mri.python.ops import linalg_ops
 from tensorflow_mri.python.ops import math_ops
 from tensorflow_mri.python.util import api_util
@@ -41,16 +42,40 @@ class ConvexFunction():
   Subclasses should implement the `_call` and `_prox` methods to define the
   forward pass and the proximal mapping, respectively. Gradients are
   provided by TensorFlow's automatic differentiation feature.
+
+  This class exposes three properties to get static shape information:
+
+  * `shape`: The static shape. Calls `_shape`.
+  * `domain_dimension`: The static domain dimension, equal to `shape[-1]`.
+  * `batch_shape`: The static batch shape, equal to `shape[:-1]`.
+
+  Additionally there are three equivalent methods to get dynamic shape
+  information:
+
+  * `shape_tensor`: The dynamic shape. Calls `_shape_tensor`.
+  * `domain_dimension_tensor`: The dynamic domain dimension, equal to
+    `shape_tensor()[-1]`.
+  * `batch_shape_tensor`: The dynamic batch shape, equal to
+    `shape_tensor()[:-1]`.
+
+  Subclasses may implement the methods `_shape` and `_shape_tensor` to provide
+  custom static and dynamic shape information, respectively.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
     """Initialize this `ConvexFunction`."""
-    if isinstance(ndim, tf.compat.v1.Dimension):
-      ndim = ndim.value
-    self._ndim = check_util.validate_rank(ndim, 'ndim', accept_none=True)
+    domain_dimension = deprecation.deprecated_argument_lookup(
+        'domain_dimension', domain_dimension, 'ndim', ndim)
+    if isinstance(domain_dimension, tf.compat.v1.Dimension):
+      domain_dimension = domain_dimension.value
+    self._domain_dimension = check_util.validate_rank(
+        domain_dimension, 'domain_dimension', accept_none=True)
     self._dtype = tf.dtypes.as_dtype(dtype or tf.dtypes.float32)
     self._name = name or type(self).__name__
     if scale is None:
@@ -106,21 +131,6 @@ def conj(self, name=None):
     with self._name_scope(name or "conj"):
       return self._conj()
 
-  def ndim_tensor(self, name=None):
-    """Returns the number of dimensions of this `ConvexFunction`.
-
-    Args:
-      name: A name for this operation (optional).
-
-    Returns:
-      A scalar integer `tf.Tensor`.
-    """
-    with self._name_scope(name or "ndim_tensor"):
-      # Prefer to use statically defined ndim if available.
-      if isinstance(self.ndim, int):
-        return tf.constant(self.ndim, dtype=tf.int32)
-      return self._ndim_tensor()
-
   def shape_tensor(self, name=None):
     """Returns the dynamic shape of this `ConvexFunction`.
 
@@ -137,9 +147,41 @@ def shape_tensor(self, name=None):
             self.shape.as_list(), name="shape")
       return self._shape_tensor()
 
+  def domain_dimension_tensor(self, name=None):
+    """Returns the dynamic domain dimension of this `ConvexFunction`.
+
+    Subclasses get this for free once they implement `_shape_tensor`.
+
+    Args:
+      name: A name for this operation (optional).
+
+    Returns:
+      A scalar integer `tf.Tensor`.
+    """
+    with self._name_scope(name or "domain_dimension_tensor"):
+      # Prefer to use statically defined domain_dimension if available.
+      if isinstance(self.domain_dimension, int):
+        return tf.constant(self.domain_dimension, dtype=tf.int32)
+      return self.shape_tensor()[-1]
+
+  @deprecation.deprecated(
+      '2022-08-07', 'Use `ConvexFunction.domain_dimension_tensor` instead.')
+  def ndim_tensor(self, name=None):  # pylint: disable=unused-argument
+    """Returns the dynamic domain dimension of this `ConvexFunction`.
+
+    Args:
+      name: A name for this operation (optional).
+
+    Returns:
+      A scalar integer `tf.Tensor`.
+    """
+    return self.domain_dimension_tensor()
+
   def batch_shape_tensor(self, name=None):
     """Returns the dynamic batch shape of this `ConvexFunction`.
 
+    Subclasses get this for free once they implement `_shape_tensor`.
+
     Args:
       name: A name for this operation (optional).
 
@@ -147,7 +189,11 @@ def batch_shape_tensor(self, name=None):
       A 1D integer `tf.Tensor`.
     """
     with self._name_scope(name or "batch_shape_tensor"):
-      return self._batch_shape_tensor()
+      # Prefer to use statically defined shape if available.
+      if self.batch_shape.is_fully_defined():
+        return tensor_util.convert_shape_to_tensor(
+            self.batch_shape.as_list(), name="shape")
+      return self.shape_tensor()[:-1]
 
   @abc.abstractmethod
   def _call(self, x):
@@ -165,21 +211,12 @@ def _conj(self):
   def _shape(self):
     """Returns the static shape of this `ConvexFunction`.
 
-    Defaults to `[ndim]`. Subclasses should override this method if they
-    have any batch dimensions.
+    Has a default implementation based on the domain dimension.
 
     Returns:
       A `tf.TensorShape`.
     """
-    return tf.TensorShape([self.ndim])
-
-  def _ndim_tensor(self):
-    """Returns the dynamic number of dimensions of this `ConvexFunction`.
-
-    Returns:
-      A scalar integer `tf.Tensor`.
-    """
-    raise NotImplementedError("`_ndim_tensor` is not implemented.")
+    return tf.TensorShape([self._domain_dimension])
 
   def _shape_tensor(self):
     """Returns the dynamic shape of this `ConvexFunction`.
@@ -189,32 +226,28 @@ def _shape_tensor(self):
     """
     raise NotImplementedError("`_shape_tensor` is not implemented.")
 
-  def _batch_shape_tensor(self):
-    """Returns the dynamic batch shape of this `ConvexFunction`.
-
-    Returns:
-      A 1D integer `tf.Tensor`.
-    """
-    if self.batch_shape.is_fully_defined():
-      return tensor_util.convert_shape_to_tensor(
-          self.batch_shape.as_list(), name="batch_shape")
-    return self.shape_tensor()[:-1]
-
   @property
   def scale(self):
     """The scaling factor."""
     return self._scale
 
-  @property
-  def ndim(self):
-    """The number of dimensions of this `ConvexFunction`."""
-    return self._ndim
-
   @property
   def shape(self):
     """The static shape of this `ConvexFunction`."""
     return self._shape()
 
+  @property
+  def domain_dimension(self):
+    """The static domain dimension of this `ConvexFunction`."""
+    return self.shape[-1]
+
+  @property
+  @deprecation.deprecated(
+      '2022-08-07', 'Use `ConvexFunction.domain_dimension` instead.')
+  def ndim(self):
+    """The number of dimensions of this `ConvexFunction`."""
+    return self.domain_dimension
+
   @property
   def batch_shape(self):
     """The static batch shape of this `ConvexFunction`."""
@@ -240,26 +273,26 @@ def _name_scope(self, name=None):
       yield scope
 
   def _check_input_shape(self, arg):  # pylint: disable=missing-param-doc
-    """Check that arg.shape[-1] is compatible with self.ndim."""
+    """Check that arg.shape[-1] is compatible with self.domain_dimension."""
     if arg.shape.rank is None:
       raise ValueError(
-          "Expected argument to have known rank, but found: %s in tensor %s" %
-          (arg.shape.rank, arg))
+          "Expected argument to have known rank, but found: {} "
+          "in tensor {}".format(arg.shape.rank, arg))
     if arg.shape.rank < 1:
       raise ValueError(
-          "Expected argument to have rank >= 1, but found: %s in tensor %s" %
-          (arg.shape.rank, arg))
-    if not arg.shape[-1:].is_compatible_with([self.ndim]):
+          "Expected argument to have rank >= 1, but found: {} "
+          "in tensor {}".format(arg.shape.rank, arg))
+    if not arg.shape[-1:].is_compatible_with([self.domain_dimension]):
       raise ValueError(
-          "Expected argument to have last dimension %d, but found: %d in "
-          "tensor %s" % (self.ndim, arg.shape[-1], arg))
+          "Expected argument to have last dimension {}, but found: {} in "
+          "tensor {}".format(self.domain_dimension, arg.shape[-1], arg))
 
   def _check_input_dtype(self, arg):
     """Check that arg.dtype == self.dtype."""
     if arg.dtype.base_dtype != self.dtype:
       raise TypeError(
-          "Expected argument to have dtype %s, but found: %s in tensor %s" %
-          (self.dtype, arg.dtype, arg))
+          "Expected argument to have dtype {}, but found: {} "
+          "in tensor {}".format(self.dtype, arg.dtype, arg))
 
 
 @api_util.export("convex.ConvexFunctionAffineMappingComposition")
@@ -282,9 +315,10 @@ def __init__(self,
                constant=None,
                scale=None,
                name=None):
-    ndim = operator.domain_dimension if operator is not None else function.ndim
+    domain_dimension = (operator.domain_dimension if operator is not None
+                        else function.domain_dimension)
     super().__init__(scale=scale,
-                     ndim=ndim,
+                     domain_dimension=domain_dimension,
                      dtype=function.dtype,
                      name=name)
     self._function = function
@@ -352,19 +386,23 @@ class ConvexFunctionIndicatorBall(ConvexFunction):  # pylint: disable=abstract-m
     order: A `float`. The order of the norm. Supported values are `1`, `2`,
       `np.inf`.
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.dtypes.DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                order=2,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
     super().__init__(scale=scale,
+                     domain_dimension=domain_dimension,
                      ndim=ndim,
                      dtype=dtype,
                      name=name)
@@ -386,7 +424,7 @@ def _conj(self):
     return ConvexFunctionNorm(
         order=_conjugate_exponent(self._order),
         scale=self._scale,
-        ndim=self.ndim,
+        domain_dimension=self.domain_dimension,
         dtype=self.dtype,
         name=f"{self.name}_conj")
 
@@ -397,8 +435,8 @@ class ConvexFunctionIndicatorL1Ball(ConvexFunctionIndicatorBall):  # pylint: dis
 
   Args:
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.dtypes.DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -407,12 +445,16 @@ class ConvexFunctionIndicatorL1Ball(ConvexFunctionIndicatorBall):  # pylint: dis
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(order=1, scale=scale, ndim=ndim, dtype=dtype, name=name)
+    super().__init__(order=1, scale=scale, domain_dimension=domain_dimension,
+                     ndim=ndim, dtype=dtype, name=name)
 
 
 @api_util.export("convex.ConvexFunctionIndicatorL2Ball")
@@ -421,8 +463,8 @@ class ConvexFunctionIndicatorL2Ball(ConvexFunctionIndicatorBall):  # pylint: dis
 
   Args:
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.dtypes.DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -431,12 +473,16 @@ class ConvexFunctionIndicatorL2Ball(ConvexFunctionIndicatorBall):  # pylint: dis
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(order=2, scale=scale, ndim=ndim, dtype=dtype, name=name)
+    super().__init__(order=2, scale=scale, domain_dimension=domain_dimension,
+                     ndim=ndim, dtype=dtype, name=name)
 
 
 @api_util.export("convex.ConvexFunctionNorm")
@@ -447,8 +493,8 @@ class ConvexFunctionNorm(ConvexFunction):  # pylint: disable=abstract-method
     order: A `float`. The order of the norm. Supported values are `1`, `2`,
       `np.inf`.
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.dtypes.DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -457,13 +503,17 @@ class ConvexFunctionNorm(ConvexFunction):  # pylint: disable=abstract-method
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                order,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(scale=scale, ndim=ndim, dtype=dtype, name=name)
+    super().__init__(scale=scale, domain_dimension=domain_dimension,
+                     ndim=ndim, dtype=dtype, name=name)
     self._order = check_util.validate_enum(order, [1, 2, np.inf], name='order')
 
   def _call(self, x):
@@ -487,7 +537,7 @@ def _conj(self):
     return ConvexFunctionIndicatorBall(
         order=_conjugate_exponent(self._order),
         scale=self._scale,
-        ndim=self.ndim,
+        domain_dimension=self.domain_dimension,
         dtype=self.dtype,
         name=f"{self.name}_conj")
 
@@ -498,8 +548,8 @@ class ConvexFunctionL1Norm(ConvexFunctionNorm):  # pylint: disable=abstract-meth
 
   Args:
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.dtypes.DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -508,12 +558,16 @@ class ConvexFunctionL1Norm(ConvexFunctionNorm):  # pylint: disable=abstract-meth
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(order=1, scale=scale, ndim=ndim, dtype=dtype, name=name)
+    super().__init__(order=1, scale=scale, domain_dimension=domain_dimension,
+                     ndim=ndim, dtype=dtype, name=name)
 
 
 @api_util.export("convex.ConvexFunctionL2Norm")
@@ -522,8 +576,8 @@ class ConvexFunctionL2Norm(ConvexFunctionNorm):  # pylint: disable=abstract-meth
 
   Args:
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `string` or `DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -532,12 +586,16 @@ class ConvexFunctionL2Norm(ConvexFunctionNorm):  # pylint: disable=abstract-meth
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(order=2, scale=scale, ndim=ndim, dtype=dtype, name=name)
+    super().__init__(order=2, scale=scale, domain_dimension=domain_dimension,
+                     ndim=ndim, dtype=dtype, name=name)
 
 
 @api_util.export("convex.ConvexFunctionL2NormSquared")
@@ -546,8 +604,8 @@ class ConvexFunctionL2NormSquared(ConvexFunction):  # pylint: disable=abstract-m
 
   Args:
     scale: A `float`. A scaling factor. Defaults to 1.0.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `string` or `DType`. The type of this `ConvexFunction`. Defaults to
       `tf.dtypes.float32`.
     name: A name for this `ConvexFunction`.
@@ -556,12 +614,16 @@ class ConvexFunctionL2NormSquared(ConvexFunction):  # pylint: disable=abstract-m
     .. [1] Parikh, N., & Boyd, S. (2014). Proximal algorithms. Foundations and
       Trends in optimization, 1(3), 127-239.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=None,
                name=None):
-    super().__init__(ndim=ndim, dtype=dtype, name=name)
+    super().__init__(domain_dimension=domain_dimension, ndim=ndim,
+                     dtype=dtype, name=name)
     self._scale = scale
 
   def _call(self, x):
@@ -589,22 +651,27 @@ class ConvexFunctionTikhonov(ConvexFunctionAffineMappingComposition):  # pylint:
       Defaults to the identity operator.
     prior: A `tf.Tensor`. The prior estimate :math:`x_0`. Defaults to 0.
     scale: A `float`. The scaling factor.
-    ndim: An `int`. The dimensionality of the domain of this `ConvexFunction`.
-      Defaults to `None`.
+    domain_dimension: An `int`. The dimension of the domain. Defaults to `None`.
+    ndim: An `int`. The dimension of the domain. Defaults to `None`.
     dtype: A `tf.DType`. The dtype of the inputs. Defaults to `float32`.
     name: A name for this `ConvexFunction`.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_dimension` instead.', ('ndim', None))
   def __init__(self,
                transform=None,
                prior=None,
                scale=None,
+               domain_dimension=None,
                ndim=None,
                dtype=tf.float32,
                name=None):
-    if ndim is None and transform is not None:
-      ndim = transform.range_dimension
+    domain_dimension = deprecation.deprecated_argument_lookup(
+        'domain_dimension', domain_dimension, 'ndim', ndim)
+    if domain_dimension is None and transform is not None:
+      domain_dimension = transform.range_dimension
     function = ConvexFunctionL2NormSquared(scale=scale,
-                                           ndim=ndim,
+                                           domain_dimension=domain_dimension,
                                            dtype=dtype)
     # Stored only for external access. Not actually used for computation.
     self._transform = transform
@@ -640,20 +707,27 @@ class ConvexFunctionTotalVariation(ConvexFunctionLinearOperatorComposition):  #
 
   Args:
     scale: A `float`. A scaling factor.
-    ndim: An `int`, or a list of `ints`. The dimensionality of the domain.
-      The domain may have multiple axes.
+    domain_shape: A `tf.TensorShape`. The shape of the domain. Defaults to
+      `None`. The domain of this `ConvexFunction` may have multiple axes.
+    ndim: A `tf.TensorShape`. The shape of the domain. Defaults to
+      `None`. The domain of this `ConvexFunction` may have multiple axes.
     axis: An `int` or a list of `ints`. The axes along which to compute the
       total variation. Defaults to -1.
     dtype: A `tf.DType`. The dtype of the inputs.
     name: A name for this `ConvexFunction`.
   """
+  @deprecation.deprecated_args(
+      '2022-08-07', 'Use argument `domain_shape` instead.', ('ndim', None))
   def __init__(self,
                scale=None,
+               domain_shape=None,
                ndim=None,
                axis=-1,
                dtype=tf.float32,
                name=None):
-    domain_shape = tf.TensorShape(ndim)
+    domain_shape = deprecation.deprecated_argument_lookup(
+        'domain_shape', domain_shape, 'ndim', ndim)
+    domain_shape = tf.TensorShape(domain_shape)
     axis = check_util.validate_axis(axis, domain_shape.rank,
                                     max_length=domain_shape.rank,
                                     canonicalize="negative")
@@ -661,7 +735,7 @@ def __init__(self,
         domain_shape, axis=ax, dtype=dtype) for ax in axis]
     operator = linalg_ext.LinearOperatorVerticalStack(operators)
     function = ConvexFunctionL1Norm(scale=scale,
-                                    ndim=operator.range_dimension,
+                                    domain_dimension=operator.range_dimension,
                                     dtype=dtype)
     super().__init__(function,
                      operator=operator,
@@ -692,7 +766,7 @@ def __init__(self,
                scale=None,
                name=None):
     super().__init__(scale=scale,
-                     ndim=quadratic_coefficient.shape[-1],
+                     domain_dimension=quadratic_coefficient.shape[-1],
                      dtype=quadratic_coefficient.dtype,
                      name=name)
     self._quadratic_coefficient = quadratic_coefficient
@@ -747,10 +821,10 @@ def _validate_linear_coefficient(self, coef):  # pylint: disable=missing-param-d
       raise ValueError(
           "Expected linear coefficient to have rank >= 1, but found: %s in "
           "tensor %s" % (coef.shape.rank, coef))
-    if not coef.shape[-1:].is_compatible_with([self.ndim]):
+    if not coef.shape[-1:].is_compatible_with([self._domain_dimension]):
       raise ValueError(
           "Expected linear coefficient to have last dimension %d, but found: "
-          "%d in tensor %s" % (self.ndim, coef.shape[-1], coef))
+          "%d in tensor %s" % (self._domain_dimension, coef.shape[-1], coef))
     if coef.dtype != self.dtype:
       raise ValueError(
           "Expected linear coefficient to have dtype %s, but found: %s in "
@@ -771,7 +845,7 @@ def _shape(self):
         self._quadratic_coefficient.shape[:-2],
         self._linear_coefficient.shape[:-1],
         self._constant_coefficient.shape)
-    return batch_shape.concatenate(tf.TensorShape([self.ndim]))
+    return batch_shape.concatenate(tf.TensorShape([self._domain_dimension]))
 
   def _shape_tensor(self):
     """Returns the dynamic shape of this `ConvexFunction`."""
@@ -779,7 +853,7 @@ def _shape_tensor(self):
         tensor_util.object_shape(self._quadratic_coefficient)[:-2],
         tf.shape(self._linear_coefficient)[:-1],
         tf.shape(self._constant_coefficient))
-    return tf.concat([batch_shape_tensor, [self.ndim_tensor()]], 0)
+    return tf.concat([batch_shape_tensor, [self._domain_dimension]], 0)
 
   @property
   def quadratic_coefficient(self):
diff --git a/tensorflow_mri/python/ops/optimizer_ops.py b/tensorflow_mri/python/ops/optimizer_ops.py
index f349c802..b497a0c4 100644
--- a/tensorflow_mri/python/ops/optimizer_ops.py
+++ b/tensorflow_mri/python/ops/optimizer_ops.py
@@ -288,10 +288,10 @@ def admm_minimize(function_f,
 
     # Infer the dimensionality of the primal variables x, z from the
     # dimensionality of the domains of f and g.
-    x_ndim_static = function_f.ndim
-    z_ndim_static = function_g.ndim
-    x_ndim = prefer_static.ndim(function_f)
-    z_ndim = prefer_static.ndim(function_g)
+    x_ndim_static = function_f.domain_dimension
+    z_ndim_static = function_g.domain_dimension
+    x_ndim = prefer_static.domain_dimension(function_f)
+    z_ndim = prefer_static.domain_dimension(function_g)
 
     # Provide default values for A and B.
     if operator_a is None:
diff --git a/tensorflow_mri/python/util/api_util.py b/tensorflow_mri/python/util/api_util.py
index a2ff24dc..7430dabd 100644
--- a/tensorflow_mri/python/util/api_util.py
+++ b/tensorflow_mri/python/util/api_util.py
@@ -22,11 +22,12 @@
 
 _API_ATTR = '_api_names'
 
-_NAMESPACES = [
+_SUBMODULE_NAMES = [
     'callbacks',
     'coils',
     'convex',
     'image',
+    'initializers',
     'io',
     'math',
     'layers',
@@ -41,11 +42,12 @@
     'summary'
 ]
 
-_NAMESPACE_DOCSTRINGS = {
+_SUBMODULE_DOCSTRINGS = {
     'callbacks': "Keras callbacks.",
     'coils': "Parallel imaging operations.",
     'convex': "Convex optimization operations.",
     'image': "Image processing operations.",
+    'initializers': "Keras initializers.",
     'io': "Input/output operations.",
     'layers': "Keras layers.",
     'linalg': "Linear algebra operations.",
@@ -66,9 +68,9 @@ def get_api_symbols():
   return _API_SYMBOLS
 
 
-def get_namespaces():
-  """Returns a list of TFMRI namespaces."""
-  return _NAMESPACES
+def get_submodule_names():
+  """Returns a list of TFMRI submodule names."""
+  return _SUBMODULE_NAMES
 
 
 def get_symbol_from_name(name):
@@ -83,6 +85,34 @@ def get_symbol_from_name(name):
   return _API_SYMBOLS.get(name)
 
 
+def get_symbols_in_submodule(name):
+  """Returns the symbols in the given submodule.
+
+  Args:
+    name: Name of the submodule.
+
+  Returns:
+    A dict containing the API symbols in the submodule.
+  """
+  symbols = {}
+  for k, v in _API_SYMBOLS.items():
+    if k.startswith(name):
+      symbols[k] = v
+  return symbols
+
+
+def get_docstring_for_submodule(name):
+  """Returns the docstring for the given submodule.
+
+  Args:
+    name: Name of the submodule.
+
+  Returns:
+    The docstring for the submodule.
+  """
+  return _SUBMODULE_DOCSTRINGS[name]
+
+
 def get_canonical_name_for_symbol(symbol):
   """Get canonical name for the API symbol.
 
@@ -131,7 +161,7 @@ def decorator(symbol):
         raise ValueError(f"Invalid API name: {name}")
       # API namespace must be one of the supported ones.
       namespace, _ = name.split('.')
-      if namespace not in _NAMESPACES:
+      if namespace not in _SUBMODULE_NAMES:
         raise ValueError(f"Invalid API namespace: {namespace}")
       # API name must be unique.
       if name in _API_SYMBOLS:
@@ -179,5 +209,5 @@ def import_namespace(namespace):
   module = importlib.util.module_from_spec(spec)
   sys.modules[spec.name] = module
   spec.loader.exec_module(module)
-  module.__doc__ = _NAMESPACE_DOCSTRINGS[namespace]
+  module.__doc__ = _SUBMODULE_DOCSTRINGS[namespace]
   return module
diff --git a/tensorflow_mri/python/util/deprecation.py b/tensorflow_mri/python/util/deprecation.py
index 4044b554..8b7ff3ca 100755
--- a/tensorflow_mri/python/util/deprecation.py
+++ b/tensorflow_mri/python/util/deprecation.py
@@ -16,4 +16,9 @@
 
 from tensorflow.python.util import deprecation
 
+deprecated_alias = deprecation.deprecated_alias
 deprecated = deprecation.deprecated
+deprecated_endpoints = deprecation.deprecated_endpoints
+deprecated_args = deprecation.deprecated_args
+deprecated_arg_values = deprecation.deprecated_arg_values
+deprecated_argument_lookup = deprecation.deprecated_argument_lookup
diff --git a/tensorflow_mri/python/util/layer_util.py b/tensorflow_mri/python/util/layer_util.py
index a0cf925c..d311af9a 100644
--- a/tensorflow_mri/python/util/layer_util.py
+++ b/tensorflow_mri/python/util/layer_util.py
@@ -16,6 +16,8 @@
 
 import tensorflow as tf
 
+from tensorflow_mri.python.layers import convolutional
+
 
 def get_nd_layer(name, rank):
   """Get an N-D layer object.
@@ -41,9 +43,9 @@ def get_nd_layer(name, rank):
     ('AveragePooling', 1): tf.keras.layers.AveragePooling1D,
     ('AveragePooling', 2): tf.keras.layers.AveragePooling2D,
     ('AveragePooling', 3): tf.keras.layers.AveragePooling3D,
-    ('Conv', 1): tf.keras.layers.Conv1D,
-    ('Conv', 2): tf.keras.layers.Conv2D,
-    ('Conv', 3): tf.keras.layers.Conv3D,
+    ('Conv', 1): convolutional.Conv1D,
+    ('Conv', 2): convolutional.Conv2D,
+    ('Conv', 3): convolutional.Conv3D,
     ('ConvLSTM', 1): tf.keras.layers.ConvLSTM1D,
     ('ConvLSTM', 2): tf.keras.layers.ConvLSTM2D,
     ('ConvLSTM', 3): tf.keras.layers.ConvLSTM3D,
diff --git a/tensorflow_mri/python/util/prefer_static.py b/tensorflow_mri/python/util/prefer_static.py
index ef694e60..b79a619d 100644
--- a/tensorflow_mri/python/util/prefer_static.py
+++ b/tensorflow_mri/python/util/prefer_static.py
@@ -50,30 +50,33 @@ def batch_shape(operator):
   return operator.batch_shape_tensor()
 
 
-def ndim(operator):
-  """Returns the dimensionality of an operator.
-
-  Returns the static dimensionality of the operator if known, otherwise
-  returns the dynamic dimensionality.
+def domain_dimension(operator):
+  """Retrieves the domain dimension of an operator.
 
   Args:
-    operator: A `tfmri.convex.ConvexFunction`.
+    operator: A `tf.linalg.LinearOperator` or a `tfmri.convex.ConvexFunction`.
 
   Returns:
-    An int or scalar integer `tf.Tensor` representing the dimensionality of the
+    An int or scalar integer `tf.Tensor` representing the range dimension of the
     operator.
 
   Raises:
-    ValueError: If `operator` is not a `tfmri.convex.ConvexFunction`.
+    ValueError: If `operator` is not a `tf.linalg.LinearOperator` or a
+    `tfmri.convex.ConvexFunction`.
   """
-  if not isinstance(operator, convex_ops.ConvexFunction):
-    raise ValueError(f"Input must be a `tfmri.convex.ConvexFunction`, "
+  if not isinstance(operator, (tf.linalg.LinearOperator,
+                               convex_ops.ConvexFunction)):
+    raise ValueError(f"Input must be a `tf.linalg.LinearOperator` or a "
+                     f"`tfmri.convex.ConvexFunction`, "
                      f"but got: {type(operator)}")
 
-  if operator.ndim is not None:
-    return operator.ndim
+  dimension = operator.domain_dimension
+  if isinstance(dimension, tf.compat.v1.Dimension):
+    dimension = dimension.value
+  if dimension is not None:
+    return dimension
 
-  return operator.ndim_tensor()
+  return operator.domain_dimension_tensor()
 
 
 def range_dimension(operator):
diff --git a/tools/assets/tfmri_icon.svg b/tools/assets/tfmri_icon.svg
index 8864f6e5..1206fd3f 100644
--- a/tools/assets/tfmri_icon.svg
+++ b/tools/assets/tfmri_icon.svg
@@ -1,1 +1,248 @@
-<svg width="260" height="259" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" overflow="hidden"><defs><clipPath id="clip0"><rect x="1264" y="1238" width="260" height="259"/></clipPath><linearGradient x1="1264" y1="1367.5" x2="1524" y2="1367.5" gradientUnits="userSpaceOnUse" spreadMethod="reflect" id="fill1"><stop offset="0" stop-color="#128091"/><stop offset="0.0227273" stop-color="#128192"/><stop offset="0.0454545" stop-color="#128394"/><stop offset="0.0681818" stop-color="#128496"/><stop offset="0.0909091" stop-color="#128698"/><stop offset="0.113636" stop-color="#13889A"/><stop offset="0.136364" stop-color="#13899B"/><stop offset="0.159091" stop-color="#138B9D"/><stop offset="0.181818" stop-color="#138C9F"/><stop offset="0.204545" stop-color="#148DA0"/><stop offset="0.227273" stop-color="#148FA2"/><stop offset="0.25" stop-color="#1490A3"/><stop offset="0.272727" stop-color="#1491A5"/><stop offset="0.295455" stop-color="#1493A6"/><stop offset="0.318182" stop-color="#1594A8"/><stop offset="0.340909" stop-color="#1595A9"/><stop offset="0.363636" stop-color="#1596AA"/><stop offset="0.386364" stop-color="#1597AB"/><stop offset="0.409091" stop-color="#1599AD"/><stop offset="0.431818" stop-color="#159AAE"/><stop offset="0.454545" stop-color="#169BAF"/><stop offset="0.477273" stop-color="#169CB0"/><stop offset="0.5" stop-color="#169DB1"/><stop offset="0.522727" stop-color="#169EB2"/><stop offset="0.545455" stop-color="#169EB3"/><stop offset="0.568182" stop-color="#169FB4"/><stop offset="0.590909" stop-color="#16A0B5"/><stop offset="0.613636" stop-color="#16A1B6"/><stop offset="0.636364" stop-color="#17A1B7"/><stop offset="0.659091" stop-color="#17A2B8"/><stop offset="0.681818" stop-color="#17A3B8"/><stop offset="0.704545" stop-color="#17A3B9"/><stop offset="0.727273" stop-color="#17A4BA"/><stop offset="0.75" stop-color="#17A5BA"/><stop offset="0.772727" stop-color="#17A5BB"/><stop offset="0.795455" stop-color="#17A5BB"/><stop offset="0.818182" stop-color="#17A6BC"/><stop offset="0.840909" stop-color="#17A6BC"/><stop offset="0.863636" stop-color="#17A7BC"/><stop offset="0.886364" stop-color="#17A7BD"/><stop offset="0.909091" stop-color="#17A7BD"/><stop offset="0.931818" stop-color="#17A7BD"/><stop offset="0.954545" stop-color="#17A7BD"/><stop offset="0.977273" stop-color="#17A7BD"/><stop offset="1" stop-color="#18A8BE"/></linearGradient></defs><g clip-path="url(#clip0)" transform="translate(-1264 -1238)"><path d="M1394 1302.75C1358.1 1302.75 1329 1331.74 1329 1367.5 1329 1403.26 1358.1 1432.25 1394 1432.25 1429.9 1432.25 1459 1403.26 1459 1367.5 1459 1331.74 1429.9 1302.75 1394 1302.75ZM1394 1289.8C1437.08 1289.8 1472 1324.59 1472 1367.5 1472 1410.41 1437.08 1445.2 1394 1445.2 1350.92 1445.2 1316 1410.41 1316 1367.5 1316 1324.59 1350.92 1289.8 1394 1289.8ZM1394 1276.85C1343.74 1276.85 1303 1317.44 1303 1367.5 1303 1417.56 1343.74 1458.15 1394 1458.15 1444.26 1458.15 1485 1417.56 1485 1367.5 1485 1317.44 1444.26 1276.85 1394 1276.85ZM1394 1238C1465.8 1238 1524 1295.98 1524 1367.5 1524 1439.02 1465.8 1497 1394 1497 1322.2 1497 1264 1439.02 1264 1367.5 1264 1295.98 1322.2 1238 1394 1238Z" fill="url(#fill1)" fill-rule="evenodd"/></g></svg>
\ No newline at end of file
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<svg
+   width="260"
+   height="260"
+   overflow="hidden"
+   version="1.1"
+   id="svg819"
+   sodipodi:docname="tfmri_icon.svg"
+   inkscape:version="1.1.2 (1:1.1+202202050942+0a00cf5339)"
+   inkscape:export-filename="/home/jmt/dev/mrphys/tensorflow-mri/tools/assets/tfmri_icon_128px.png"
+   inkscape:export-xdpi="47.259998"
+   inkscape:export-ydpi="47.259998"
+   xmlns:inkscape="http://www.inkscape.org/namespaces/inkscape"
+   xmlns:sodipodi="http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd"
+   xmlns="http://www.w3.org/2000/svg"
+   xmlns:svg="http://www.w3.org/2000/svg">
+  <sodipodi:namedview
+     id="namedview821"
+     pagecolor="#ffffff"
+     bordercolor="#666666"
+     borderopacity="1.0"
+     inkscape:pageshadow="2"
+     inkscape:pageopacity="0.0"
+     inkscape:pagecheckerboard="0"
+     showgrid="false"
+     height="260px"
+     inkscape:zoom="3.3706564"
+     inkscape:cx="95.085338"
+     inkscape:cy="129.5"
+     inkscape:window-width="2492"
+     inkscape:window-height="1376"
+     inkscape:window-x="68"
+     inkscape:window-y="27"
+     inkscape:window-maximized="1"
+     inkscape:current-layer="svg819" />
+  <defs
+     id="defs813">
+    <clipPath
+       id="clip0">
+      <rect
+         x="1264"
+         y="1238"
+         width="260"
+         height="259"
+         id="rect719" />
+    </clipPath>
+    <linearGradient
+       x1="1264"
+       y1="1367.5"
+       x2="1524"
+       y2="1367.5"
+       gradientUnits="userSpaceOnUse"
+       spreadMethod="reflect"
+       id="fill1">
+      <stop
+         offset="0"
+         stop-color="#128091"
+         id="stop722" />
+      <stop
+         offset="0.0227273"
+         stop-color="#128192"
+         id="stop724" />
+      <stop
+         offset="0.0454545"
+         stop-color="#128394"
+         id="stop726" />
+      <stop
+         offset="0.0681818"
+         stop-color="#128496"
+         id="stop728" />
+      <stop
+         offset="0.0909091"
+         stop-color="#128698"
+         id="stop730" />
+      <stop
+         offset="0.113636"
+         stop-color="#13889A"
+         id="stop732" />
+      <stop
+         offset="0.136364"
+         stop-color="#13899B"
+         id="stop734" />
+      <stop
+         offset="0.159091"
+         stop-color="#138B9D"
+         id="stop736" />
+      <stop
+         offset="0.181818"
+         stop-color="#138C9F"
+         id="stop738" />
+      <stop
+         offset="0.204545"
+         stop-color="#148DA0"
+         id="stop740" />
+      <stop
+         offset="0.227273"
+         stop-color="#148FA2"
+         id="stop742" />
+      <stop
+         offset="0.25"
+         stop-color="#1490A3"
+         id="stop744" />
+      <stop
+         offset="0.272727"
+         stop-color="#1491A5"
+         id="stop746" />
+      <stop
+         offset="0.295455"
+         stop-color="#1493A6"
+         id="stop748" />
+      <stop
+         offset="0.318182"
+         stop-color="#1594A8"
+         id="stop750" />
+      <stop
+         offset="0.340909"
+         stop-color="#1595A9"
+         id="stop752" />
+      <stop
+         offset="0.363636"
+         stop-color="#1596AA"
+         id="stop754" />
+      <stop
+         offset="0.386364"
+         stop-color="#1597AB"
+         id="stop756" />
+      <stop
+         offset="0.409091"
+         stop-color="#1599AD"
+         id="stop758" />
+      <stop
+         offset="0.431818"
+         stop-color="#159AAE"
+         id="stop760" />
+      <stop
+         offset="0.454545"
+         stop-color="#169BAF"
+         id="stop762" />
+      <stop
+         offset="0.477273"
+         stop-color="#169CB0"
+         id="stop764" />
+      <stop
+         offset="0.5"
+         stop-color="#169DB1"
+         id="stop766" />
+      <stop
+         offset="0.522727"
+         stop-color="#169EB2"
+         id="stop768" />
+      <stop
+         offset="0.545455"
+         stop-color="#169EB3"
+         id="stop770" />
+      <stop
+         offset="0.568182"
+         stop-color="#169FB4"
+         id="stop772" />
+      <stop
+         offset="0.590909"
+         stop-color="#16A0B5"
+         id="stop774" />
+      <stop
+         offset="0.613636"
+         stop-color="#16A1B6"
+         id="stop776" />
+      <stop
+         offset="0.636364"
+         stop-color="#17A1B7"
+         id="stop778" />
+      <stop
+         offset="0.659091"
+         stop-color="#17A2B8"
+         id="stop780" />
+      <stop
+         offset="0.681818"
+         stop-color="#17A3B8"
+         id="stop782" />
+      <stop
+         offset="0.704545"
+         stop-color="#17A3B9"
+         id="stop784" />
+      <stop
+         offset="0.727273"
+         stop-color="#17A4BA"
+         id="stop786" />
+      <stop
+         offset="0.75"
+         stop-color="#17A5BA"
+         id="stop788" />
+      <stop
+         offset="0.772727"
+         stop-color="#17A5BB"
+         id="stop790" />
+      <stop
+         offset="0.795455"
+         stop-color="#17A5BB"
+         id="stop792" />
+      <stop
+         offset="0.818182"
+         stop-color="#17A6BC"
+         id="stop794" />
+      <stop
+         offset="0.840909"
+         stop-color="#17A6BC"
+         id="stop796" />
+      <stop
+         offset="0.863636"
+         stop-color="#17A7BC"
+         id="stop798" />
+      <stop
+         offset="0.886364"
+         stop-color="#17A7BD"
+         id="stop800" />
+      <stop
+         offset="0.909091"
+         stop-color="#17A7BD"
+         id="stop802" />
+      <stop
+         offset="0.931818"
+         stop-color="#17A7BD"
+         id="stop804" />
+      <stop
+         offset="0.954545"
+         stop-color="#17A7BD"
+         id="stop806" />
+      <stop
+         offset="0.977273"
+         stop-color="#17A7BD"
+         id="stop808" />
+      <stop
+         offset="1"
+         stop-color="#18A8BE"
+         id="stop810" />
+    </linearGradient>
+  </defs>
+  <g
+     clip-path="url(#clip0)"
+     transform="translate(-1264,-1237.5)"
+     id="g817">
+    <path
+       d="m 1394,1302.75 c -35.9,0 -65,28.99 -65,64.75 0,35.76 29.1,64.75 65,64.75 35.9,0 65,-28.99 65,-64.75 0,-35.76 -29.1,-64.75 -65,-64.75 z m 0,-12.95 c 43.08,0 78,34.79 78,77.7 0,42.91 -34.92,77.7 -78,77.7 -43.08,0 -78,-34.79 -78,-77.7 0,-42.91 34.92,-77.7 78,-77.7 z m 0,-12.95 c -50.26,0 -91,40.59 -91,90.65 0,50.06 40.74,90.65 91,90.65 50.26,0 91,-40.59 91,-90.65 0,-50.06 -40.74,-90.65 -91,-90.65 z m 0,-38.85 c 71.8,0 130,57.98 130,129.5 0,71.52 -58.2,129.5 -130,129.5 -71.8,0 -130,-57.98 -130,-129.5 0,-71.52 58.2,-129.5 130,-129.5 z"
+       fill="url(#fill1)"
+       fill-rule="evenodd"
+       id="path815"
+       style="fill:url(#fill1)" />
+  </g>
+</svg>
diff --git a/tools/assets/tfmri_icon_128px.png b/tools/assets/tfmri_icon_128px.png
new file mode 100644
index 00000000..49145d3d
Binary files /dev/null and b/tools/assets/tfmri_icon_128px.png differ
diff --git a/tools/assets/tfmri_icon_192px.png b/tools/assets/tfmri_icon_192px.png
new file mode 100644
index 00000000..cf524b04
Binary files /dev/null and b/tools/assets/tfmri_icon_192px.png differ
diff --git a/tools/assets/tfmri_icon_24px.png b/tools/assets/tfmri_icon_24px.png
new file mode 100644
index 00000000..113b0f1a
Binary files /dev/null and b/tools/assets/tfmri_icon_24px.png differ
diff --git a/tools/assets/tfmri_icon_256px.png b/tools/assets/tfmri_icon_256px.png
new file mode 100644
index 00000000..de3254ed
Binary files /dev/null and b/tools/assets/tfmri_icon_256px.png differ
diff --git a/tools/assets/tfmri_icon_384px.png b/tools/assets/tfmri_icon_384px.png
new file mode 100644
index 00000000..51760056
Binary files /dev/null and b/tools/assets/tfmri_icon_384px.png differ
diff --git a/tools/assets/tfmri_icon_48px.png b/tools/assets/tfmri_icon_48px.png
new file mode 100644
index 00000000..86e9ca6d
Binary files /dev/null and b/tools/assets/tfmri_icon_48px.png differ
diff --git a/tools/assets/tfmri_icon_96px.png b/tools/assets/tfmri_icon_96px.png
new file mode 100644
index 00000000..3bea671b
Binary files /dev/null and b/tools/assets/tfmri_icon_96px.png differ
diff --git a/tools/build/create_api.py b/tools/build/create_api.py
new file mode 100644
index 00000000..a8cebd76
--- /dev/null
+++ b/tools/build/create_api.py
@@ -0,0 +1,124 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Creates the public API of TensorFlow MRI."""
+
+import inspect
+import pathlib
+import string
+import sys
+
+SCRIPT_PATH = pathlib.Path(__file__).resolve()
+BUILD_PATH = SCRIPT_PATH.parent
+ROOT_PATH = BUILD_PATH.parent.parent
+API_PATH = ROOT_PATH / 'tensorflow_mri/_api'
+INIT_PATH = ROOT_PATH / 'tensorflow_mri/__init__.py'
+
+sys.path.insert(0, str(ROOT_PATH))
+from tensorflow_mri.python.util import api_util as api_util
+
+INIT_TEMPLATE = string.Template(
+'''# This file was automatically generated by ${script_path}.
+# Do not edit.
+"""TensorFlow MRI."""
+import os as _os
+import sys as _sys
+
+from tensorflow_mri.__about__ import *
+
+# TODO(jmontalt): Remove these imports on release 1.0.0.
+from tensorflow_mri.python.ops.array_ops import *
+from tensorflow_mri.python.ops.coil_ops import *
+from tensorflow_mri.python.ops.convex_ops import *
+from tensorflow_mri.python.ops.fft_ops import *
+from tensorflow_mri.python.ops.geom_ops import *
+from tensorflow_mri.python.ops.image_ops import *
+from tensorflow_mri.python.ops.linalg_ops import *
+from tensorflow_mri.python.ops.math_ops import *
+from tensorflow_mri.python.ops.optimizer_ops import *
+from tensorflow_mri.python.ops.recon_ops import *
+from tensorflow_mri.python.ops.signal_ops import *
+from tensorflow_mri.python.ops.traj_ops import *
+
+from tensorflow_mri import python
+
+# Import submodules.
+${submodule_imports}
+
+# Make sure directory containing top level submodules is in
+# the __path__ so that "from tensorflow_mri.foo import bar" works.
+# We're using callbacks, but there's nothing special about that.
+_API_MODULE = _sys.modules[__name__].callbacks
+_tfmri_api_dir = _os.path.dirname(_os.path.dirname(_API_MODULE.__file__))
+_current_module = _sys.modules[__name__]
+
+if not hasattr(_current_module, '__path__'):
+  __path__ = [_tfmri_api_dir]
+elif _tfmri_api_dir not in __path__:
+  __path__.append(_tfmri_api_dir)
+''')
+
+
+SUBMODULE_TEMPLATE = string.Template(
+'''# This file was automatically generated by ${script_path}.
+# Do not edit.
+"""${docstring}"""
+
+${symbol_imports}
+''')
+
+SUBMODULE_IMPORT_TEMPLATE = "from tensorflow_mri._api import {submodule_name}"
+SYMBOL_IMPORT_TEMPLATE = "from {module_name} import {symbol_name} as {symbol_alias}"
+
+# Update the top-level __init__.py file.
+submodule_imports = [
+    SUBMODULE_IMPORT_TEMPLATE.format(submodule_name=submodule_name)
+    for submodule_name in api_util.get_submodule_names()]
+submodule_imports = '\n'.join(submodule_imports)
+
+init_contents = INIT_TEMPLATE.substitute(
+    script_path=SCRIPT_PATH.relative_to(ROOT_PATH),
+    submodule_imports=submodule_imports)
+
+with open(INIT_PATH, 'w') as f:
+  f.write(init_contents)
+
+# Now generate the individual submodule APIs.
+for submodule_name in api_util.get_submodule_names():
+
+  docstring = api_util.get_docstring_for_submodule(submodule_name)
+  symbols = api_util.get_symbols_in_submodule(submodule_name)
+
+  symbol_imports = []
+  for api_name, symbol in symbols.items():
+    symbol_alias = api_name.split('.')[-1]
+    module_name = inspect.getmodule(symbol).__name__
+    symbol_name = symbol.__name__
+    symbol_imports.append(
+        SYMBOL_IMPORT_TEMPLATE.format(module_name=module_name,
+                                      symbol_name=symbol_name,
+                                      symbol_alias=symbol_alias))
+  symbol_imports = '\n'.join(symbol_imports)
+
+  submodule_path = API_PATH / submodule_name
+  submodule_path.mkdir(parents=True, exist_ok=True)
+  submodule_path = submodule_path / '__init__.py'
+
+  out = SUBMODULE_TEMPLATE.substitute(
+      script_path=SCRIPT_PATH.relative_to(ROOT_PATH),
+      docstring=docstring,
+      symbol_imports=symbol_imports)
+
+  with open(submodule_path, 'w') as f:
+    f.write(out)
diff --git a/tools/docs/conf.py b/tools/docs/conf.py
index d62129ed..d04da614 100644
--- a/tools/docs/conf.py
+++ b/tools/docs/conf.py
@@ -105,7 +105,9 @@
   'dark_logo': 'tmfri_logo_dark.svg',
   'light_css_variables': {
     'color-brand-primary': '#128091',
-    'color-brand-content': '#128091'
+    'color-brand-content': '#128091',
+    'font-stack': 'Roboto, sans-serif',
+    "font-stack--monospace": "Roboto Mono, monospace"
   },
   'dark_css_variables': {
     'color-brand-primary': '#18A8BE',
@@ -113,6 +115,10 @@
   }
 }
 
+html_css_files = [
+    'https://fonts.googleapis.com/css?family=Roboto|Roboto+Mono',
+]
+
 # Additional files to copy to output directory.
 html_extra_path = ['robots.txt']
 
@@ -158,21 +164,34 @@ def linkcode_resolve(domain, info):
   # If no file, we're done. This happens for C++ ops.
   if file is None:
     return None
+  # When using TF's deprecation decorators, `getsourcefile` returns the
+  # `deprecation.py` file where the decorators are defined instead of the
+  # file where the object is defined. This should probably be fixed on the
+  # decorators themselves. For now, we just don't add the link for deprecated
+  # objects.
+  if 'deprecation' in file:
+    return None
   # Crop anything before `tensorflow_mri\python`. This path is system
   # dependent and we don't care about it.
   index = file.index('tensorflow_mri/python')
   file = file[index:]
 
-  # Get first and last line numbers.
-  lines, start = inspect.getsourcelines(obj)
-  stop = start + len(lines) - 1
-
   # Base URL.
   url = 'https://github.com/mrphys/tensorflow-mri'
   # Add version blob.
   url += '/blob/v' + release
   # Add file.
   url += '/' + file
+
+  # Try to add line numbers. This will not work when the class is defined
+  # dynamically. In that case we point to the file, but no line number.
+  try:
+    lines, start = inspect.getsourcelines(obj)
+    stop = start + len(lines) - 1
+  except OSError:
+    # Could not get source lines.
+    return url
+
   # Add line numbers.
   url += '#L' + str(start) + '-L' + str(stop)
 
@@ -229,10 +248,47 @@ def linkcode_resolve(domain, info):
 
 def process_docstring(app, what, name, obj, options, lines):  # pylint: disable=missing-param-doc,unused-argument
   """Process autodoc docstrings."""
+  # Replace Note: and Warning: by RST equivalents.
+  rst_lines = []
+  admonition_lines = None
+  for line in lines:
+    if admonition_lines is None:
+      # We are not in an admonition right now. Check if this line will start
+      # one.
+      if (line.strip().startswith('Warning:') or
+          line.strip().startswith('Note:')):
+        # This line starts an admonition.
+        label_position = line.index(':')
+        admonition_type = line[:label_position].strip().lower()
+        admonition_content = line[label_position + 1:].strip()
+        leading_whitespace = ' ' * (len(line) - len(line.lstrip()))
+        extra_indentation = '  '
+        admonition_lines = [f"{leading_whitespace}.. {admonition_type}::"]
+        admonition_lines.append(
+            leading_whitespace + extra_indentation + admonition_content)
+      else:
+        # This line does not start an admonition. It's just a regular line.
+        # Add it to the new lines.
+        rst_lines.append(line)
+    else:
+      # Check if this is the end of the admonition.
+      if line.strip() == '':
+        # Line is empty, so the end of the admonition. Add admonition and
+        # finish.
+        rst_lines.extend(admonition_lines)
+        admonition_lines = None
+      else:
+        # This is an admonition line. Add to list of admonition lines.
+        admonition_lines.append(extra_indentation + line)
+  # If we reached the end and we are still in an admonition, add it.
+  if admonition_lines is not None:
+    rst_lines.extend(admonition_lines)
+
   # Replace markdown literal markers (`) by ReST literal markers (``).
-  myst = '\n'.join(lines)
+  myst = '\n'.join(rst_lines)
   text = myst.replace('`', '``')
   text = text.replace(':math:``', ':math:`')
+
   # Correct inline code followed by word characters.
   text = CODE_LETTER_PATTERN.sub(CODE_LETTER_REPL, text)
   # Add links to some common types.
diff --git a/tools/docs/generate_documents.py b/tools/docs/create_documents.py
similarity index 94%
rename from tools/docs/generate_documents.py
rename to tools/docs/create_documents.py
index daf793a9..c3c53ede 100644
--- a/tools/docs/generate_documents.py
+++ b/tools/docs/create_documents.py
@@ -115,7 +115,7 @@ class Module:
   classes: typing.List[str] = dataclasses.field(default_factory=list)
   functions: typing.List[str] = dataclasses.field(default_factory=list)
 
-modules = {namespace: Module() for namespace in api_util.get_namespaces()}
+modules = {namespace: Module() for namespace in api_util.get_submodule_names()}
 
 for name, symbol in api_util.get_api_symbols().items():
   name = api_util.get_canonical_name_for_symbol(symbol)
@@ -142,14 +142,14 @@ class Module:
 # Write top-level API doc tfmri.rst.
 filename = os.path.join(API_DOCS_PATH, 'tfmri.rst')
 with open(filename, 'w') as f:
-  namespaces = api_util.get_namespaces()
+  namespaces = api_util.get_submodule_names()
   f.write(TFMRI_DOC_TEMPLATE.substitute(
-      namespaces='\n   '.join(sorted(namespaces))))
+      namespaces='\n    '.join(sorted(namespaces))))
 
 # Write index.rst.
 filename = os.path.join(DOCS_PATH, 'index.rst')
 with open(filename, 'w') as f:
-  namespaces = api_util.get_namespaces()
+  namespaces = api_util.get_submodule_names()
   namespaces = ['api_docs/tfmri/' + namespace for namespace in namespaces]
   f.write(INDEX_TEMPLATE.substitute(
       namespaces='\n   '.join(sorted(namespaces))))
diff --git a/tools/docs/generate_templates.py b/tools/docs/create_templates.py
similarity index 93%
rename from tools/docs/generate_templates.py
rename to tools/docs/create_templates.py
index bed3f52f..55e651a5 100644
--- a/tools/docs/generate_templates.py
+++ b/tools/docs/create_templates.py
@@ -44,16 +44,7 @@
 .. auto{{ objtype }}:: {{ objname }}
 """)
 
-NAMESPACES = api_util.get_namespaces()
-NAMESPACES += [
-    'callbacks',
-    'io',
-    'layers',
-    'linalg',
-    'losses',
-    'metrics',
-    'ops'
-]
+NAMESPACES = api_util.get_submodule_names()
 
 TEMPLATE_PATH = os.path.join(
     os.path.dirname(os.path.realpath(__file__)), '_templates')


Test patch:
diff --git a/tensorflow_mri/python/initializers/initializers_test.py b/tensorflow_mri/python/initializers/initializers_test.py
new file mode 100644
index 00000000..c98f5820
--- /dev/null
+++ b/tensorflow_mri/python/initializers/initializers_test.py
@@ -0,0 +1,261 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for Keras initializers."""
+# pylint: disable=missing-function-docstring
+
+from keras import backend
+from keras import combinations
+from keras import models
+from keras.engine import input_layer
+from keras.layers import core
+import numpy as np
+import tensorflow as tf
+import tensorflow_mri as tfmri
+
+
+def _compute_fans(shape):
+  """Computes the number of input and output units for a weight shape.
+  Args:
+    shape: Integer shape tuple or TF tensor shape.
+  Returns:
+    A tuple of integer scalars (fan_in, fan_out).
+  """
+  if len(shape) < 1:  # Just to avoid errors for constants.
+    fan_in = fan_out = 1
+  elif len(shape) == 1:
+    fan_in = fan_out = shape[0]
+  elif len(shape) == 2:
+    fan_in = shape[0]
+    fan_out = shape[1]
+  else:
+    # Assuming convolution kernels (2D, 3D, or more).
+    # kernel shape: (..., input_depth, depth)
+    receptive_field_size = 1
+    for dim in shape[:-2]:
+      receptive_field_size *= dim
+    fan_in = shape[-2] * receptive_field_size
+    fan_out = shape[-1] * receptive_field_size
+  return int(fan_in), int(fan_out)
+
+
+@combinations.generate(combinations.combine(mode=['graph', 'eager']))
+class KerasInitializersTest(tf.test.TestCase):
+  """Tests for Keras initializers."""
+  def _runner(self, init, shape, target_mean=None, target_std=None,  # pylint: disable=missing-function-docstring,unused-argument
+              target_max=None, target_min=None, dtype=None):  # pylint: disable=unused-argument
+    # The global seed is set so that we can get the same random streams between
+    # eager and graph mode when stateful op is used.
+    tf.random.set_seed(1337)
+    variable = backend.variable(init(shape, dtype), dtype)
+    output = backend.get_value(variable)
+
+    # Test output dtype.
+    self.assertDTypeEqual(output, dtype or backend.floatx())
+
+    # Test serialization (assumes deterministic behavior).
+    config = init.get_config()
+    reconstructed_init = init.__class__.from_config(config)
+
+    tf.random.set_seed(1337)
+    variable = backend.variable(reconstructed_init(shape, dtype), dtype)
+    output_2 = backend.get_value(variable)
+    self.assertAllClose(output, output_2, atol=1e-4)
+
+  def test_lecun_uniform(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(1. / fan_in)
+      self._runner(
+          tfmri.initializers.LecunUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_lecun_uniform_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, fan_out = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / (fan_in + fan_out))
+      self._runner(
+          tfmri.initializers.LecunUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_glorot_uniform(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, fan_out = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / (fan_in + fan_out))
+      self._runner(
+          tfmri.initializers.GlorotUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_glorot_uniform_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, fan_out = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / (fan_in + fan_out))
+      self._runner(
+          tfmri.initializers.GlorotUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_he_uniform(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / fan_in)
+      self._runner(
+          tfmri.initializers.HeUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_he_uniform_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / fan_in)
+      self._runner(
+          tfmri.initializers.HeUniform(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_lecun_normal(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(1. / fan_in)
+      self._runner(
+          tfmri.initializers.LecunNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_lecun_normal_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(1. / fan_in)
+      self._runner(
+          tfmri.initializers.LecunNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_glorot_normal(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, fan_out = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / (fan_in + fan_out))
+      self._runner(
+          tfmri.initializers.GlorotNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_glorot_normal_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, fan_out = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / (fan_in + fan_out))
+      self._runner(
+          tfmri.initializers.GlorotNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_he_normal(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / fan_in)
+      self._runner(
+          tfmri.initializers.HeNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std)
+
+  def test_he_normal_complex(self):
+    tensor_shape = (5, 6, 4, 2)
+    with self.cached_session():
+      fan_in, _ = _compute_fans(tensor_shape)
+      std = np.sqrt(2. / fan_in)
+      self._runner(
+          tfmri.initializers.HeNormal(seed=123),
+          tensor_shape,
+          target_mean=0.,
+          target_std=std,
+          dtype='complex64')
+
+  def test_tfmri_initializer_saving(self):
+
+    inputs = input_layer.Input((10,))
+    outputs = core.Dense(
+        1, kernel_initializer=tfmri.initializers.GlorotUniform())(inputs)
+    model = models.Model(inputs, outputs)
+    model2 = model.from_config(model.get_config())
+    self.assertIsInstance(model2.layers[1].kernel_initializer,
+                          tfmri.initializers.GlorotUniform)
+
+  def test_partition(self):
+    with self.cached_session():
+      partition_enabled_initializers = [
+          tfmri.initializers.LecunUniform(),
+          tfmri.initializers.GlorotUniform(),
+          tfmri.initializers.HeUniform()
+      ]
+      for initializer in partition_enabled_initializers:
+        got = initializer(
+            shape=(4, 2), partition_shape=(2, 2), partition_offset=(0, 0))
+        self.assertEqual(got.shape, (2, 2))
+
+      partition_forbidden_initializers = []
+      for initializer in partition_forbidden_initializers:
+        with self.assertRaisesRegex(
+            ValueError,
+            "initializer doesn't support partition-related arguments"):
+          initializer(
+              shape=(4, 2), partition_shape=(2, 2), partition_offset=(0, 0))
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensorflow_mri/python/io/twix_io_test.py b/tensorflow_mri/python/io/twix_io_test.py
new file mode 100644
index 00000000..52b42efa
--- /dev/null
+++ b/tensorflow_mri/python/io/twix_io_test.py
@@ -0,0 +1,88 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for module `twix_io`."""
+
+import tensorflow as tf
+import tensorflow_mri as tfmri
+
+from tensorflow_mri.python.util import test_util
+
+TEST_FILE = 'tests/data/gre.dat'
+
+class ParseTwixTest(test_util.TestCase):
+  """Tests for `parse_twix` function."""
+
+  def test_parse_twix(self):
+    """Tests reading from a file."""
+    contents = tf.io.read_file(TEST_FILE)
+
+    # Parse TWIX file.
+    twix = tfmri.io.parse_twix(contents)
+
+    # Basic looping.
+    kspace = []
+    for scan in twix.measurements[0].scans:
+      if not scan.header.eval_info_mask.ACQEND:
+        kspace.append([])
+        for channel in scan.channels:
+          kspace[-1].append(channel.data)
+    kspace = tf.convert_to_tensor(kspace)
+    kspace = tf.transpose(kspace, [1, 0, 2])
+
+    # Reconstruct an image.
+    image = tfmri.recon.adj(kspace, [160, 320])
+
+    # Test 40 pixels.
+    test_pixels = [
+        10274, 100414,  43817,  82566,  47161,  95817,  73454,  37660,
+        17116,  37163,  36018,  94110,  61346,  30678,   3381,   2638,
+        84365,  65707,   4120,  14389,  53742,  82352,  72599,  31465,
+        34866,  92797,  40474,  58875,  80353,  86533,  61116,  39078,
+        17864,  33563,  73549,  49704,  89011,  64354,  24272,  30771
+    ]
+
+    expected_values = [
+        -1.2701083e-06-2.3950960e-07j,  1.5321851e-06-6.0459581e-07j,
+        -2.3624499e-07-3.3690151e-07j, -1.8920543e-08+3.6719811e-07j,
+         6.7760725e-06+6.7225542e-06j,  2.0235464e-05-2.6655940e-07j,
+         1.2165106e-06-6.7168941e-07j,  1.8745079e-05-1.2247641e-05j,
+         1.6538288e-06+2.0086329e-06j, -5.8484528e-07-1.4344830e-06j,
+         9.6985887e-06-1.6889178e-05j,  7.2497940e-08+2.2277927e-06j,
+        -7.6203497e-07+2.8434135e-06j, -1.8966676e-06+7.0185899e-07j,
+         3.1993693e-06-4.0846868e-07j,  2.3145808e-06-4.0182917e-07j,
+         1.4444863e-05-1.6325985e-05j,  1.2718554e-05+2.8837078e-06j,
+        -2.5439724e-07-1.5061369e-07j,  9.1443229e-07+2.6340251e-07j,
+        -1.6596479e-07-3.6424115e-07j,  2.0008658e-05+5.8169287e-09j,
+        -4.9824121e-07+1.8819964e-07j,  1.6354710e-05+5.4614238e-06j,
+         7.8803828e-07-5.8470232e-07j,  3.4816506e-07+1.2416560e-08j,
+         1.3673757e-05-1.0114289e-05j,  2.3925327e-07+5.3502725e-07j,
+        -4.6500159e-07-4.9572600e-07j, -2.6744103e-06+9.4111897e-07j,
+        -8.6854664e-07-4.1762880e-07j, -4.3979887e-07-1.8401410e-07j,
+        -4.7728474e-07+1.3821273e-06j, -1.4986534e-07+9.7267298e-07j,
+        -1.1740372e-06-1.2283937e-06j, -1.2053533e-06+2.8251404e-06j,
+         7.2834683e-07+1.4818811e-06j, -2.4099645e-06-5.7378662e-07j,
+        -5.6130324e-07-1.1934012e-06j,  2.0544548e-09+1.2545783e-06j]
+
+    flat_image = tf.reshape(image, [-1])
+    self.assertAllEqual([2, 160, 320], tf.shape(image))
+    self.assertAllClose(expected_values, tf.gather(flat_image, test_pixels))  # pylint: disable=no-value-for-parameter
+
+    # Check protocol.
+    meas_prot = twix.measurements[0].protocol['Meas']
+    self.assertAllEqual(160, meas_prot.MEAS.sKSpace.lBaseResolution.value)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensorflow_mri/python/layers/convolutional_test.py b/tensorflow_mri/python/layers/convolutional_test.py
new file mode 100644
index 00000000..b4d81549
--- /dev/null
+++ b/tensorflow_mri/python/layers/convolutional_test.py
@@ -0,0 +1,856 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for convolutional layers."""
+# pylint: disable=missing-class-docstring,missing-function-docstring
+
+# Using some private functions. This works as of TF/Keras 2.8.0.
+from absl.testing import parameterized
+import keras
+from keras import keras_parameterized
+from keras import testing_utils
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.framework import test_util
+import tensorflow_mri as tfmri
+
+
+@keras_parameterized.run_all_keras_modes
+class Conv1DTest(keras_parameterized.TestCase):
+
+  def _run_test(self, kwargs, expected_output_shape):
+    num_samples = 2
+    stack_size = 3
+    length = 7
+
+    with self.cached_session():
+      testing_utils.layer_test(
+          tfmri.layers.Conv1D,
+          kwargs=kwargs,
+          input_shape=(num_samples, length, stack_size),
+          expected_output_shape=expected_output_shape)
+
+  def _run_test_extra_batch_dim(self, kwargs, expected_output_shape):
+    batch_shape = (2, 11)
+    stack_size = 3
+    length = 7
+
+    with self.cached_session():
+      if expected_output_shape is not None:
+        expected_output_shape = (None,) + expected_output_shape
+
+      testing_utils.layer_test(
+          tfmri.layers.Conv1D,
+          kwargs=kwargs,
+          input_shape=batch_shape + (length, stack_size),
+          expected_output_shape=expected_output_shape)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 5, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 7, 2)),
+      ('padding_same_dilation_2', {
+          'padding': 'same',
+          'dilation_rate': 2
+      }, (None, 7, 2)),
+      ('padding_same_dilation_3', {
+          'padding': 'same',
+          'dilation_rate': 3
+      }, (None, 7, 2)),
+      ('padding_causal', {
+          'padding': 'causal'
+      }, (None, 7, 2)),
+      ('strides', {
+          'strides': 2
+      }, (None, 3, 2)),
+      ('dilation_rate', {
+          'dilation_rate': 2
+      }, (None, 3, 2)),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 5, 6), True),
+  )
+  def test_conv1d(self, kwargs, expected_output_shape, requires_gpu=False):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = 3
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      self._run_test(kwargs, expected_output_shape)
+      self._run_test_extra_batch_dim(kwargs, expected_output_shape)
+
+  def test_conv1d_regularizers(self):
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_regularizer': 'l2',
+        'bias_regularizer': 'l2',
+        'activity_regularizer': 'l2',
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv1D(**kwargs)
+      layer.build((None, 5, 2))
+      self.assertEqual(len(layer.losses), 2)
+      layer(keras.backend.variable(np.ones((1, 5, 2))))
+      self.assertEqual(len(layer.losses), 3)
+
+  def test_conv1d_constraints(self):
+    k_constraint = lambda x: x
+    b_constraint = lambda x: x
+
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_constraint': k_constraint,
+        'bias_constraint': b_constraint,
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv1D(**kwargs)
+      layer.build((None, 5, 2))
+      self.assertEqual(layer.kernel.constraint, k_constraint)
+      self.assertEqual(layer.bias.constraint, b_constraint)
+
+  def test_conv1d_recreate_conv(self):
+    with self.cached_session():
+      layer = tfmri.layers.Conv1D(filters=1,
+                                  kernel_size=3,
+                                  strides=1,
+                                  dilation_rate=2,
+                                  padding='causal')
+      inpt1 = np.random.normal(size=[1, 2, 1])
+      inpt2 = np.random.normal(size=[1, 1, 1])
+      outp1_shape = layer(inpt1).shape
+      _ = layer(inpt2).shape
+      self.assertEqual(outp1_shape, layer(inpt1).shape)
+
+  def test_conv1d_recreate_conv_unknown_dims(self):
+    with self.cached_session():
+      layer = tfmri.layers.Conv1D(filters=1,
+                                  kernel_size=3,
+                                  strides=1,
+                                  dilation_rate=2,
+                                  padding='causal')
+
+      inpt1 = np.random.normal(size=[1, 9, 1]).astype(np.float32)
+      inpt2 = np.random.normal(size=[1, 2, 1]).astype(np.float32)
+      outp1_shape = layer(inpt1).shape
+
+      @tf.function(input_signature=[
+          tf.TensorSpec([1, None, 1])])
+      def fn(inpt):
+        return layer(inpt)
+
+      fn(inpt2)
+      self.assertEqual(outp1_shape, layer(inpt1).shape)
+
+  def test_conv1d_invalid_output_shapes(self):
+    kwargs = {'filters': 2, 'kernel_size': 20}
+    with self.assertRaisesRegex(
+        ValueError, r"""One of the dimensions in the output is <= 0"""):
+      layer = tfmri.layers.Conv1D(**kwargs)
+      layer.build((None, 5, 2))
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 5, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 7, 2)),
+      ('padding_same_dilation_2', {
+          'padding': 'same',
+          'dilation_rate': 2
+      }, (None, 7, 2)),
+      ('padding_same_dilation_3', {
+          'padding': 'same',
+          'dilation_rate': 3
+      }, (None, 7, 2)),
+      ('padding_causal', {
+          'padding': 'causal'
+      }, (None, 7, 2)),
+      ('strides', {
+          'strides': 2
+      }, (None, 3, 2)),
+      ('dilation_rate', {
+          'dilation_rate': 2
+      }, (None, 3, 2)),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 5, 6), True),
+      ('activation', {
+          'activation': 'exponential'
+      }, (None, 5, 2)),
+      ('regularizer', {
+          'kernel_regularizer': 'l2'
+      }, (None, 5, 2))
+  )
+  def test_conv1d_complex(
+      self, kwargs, expected_output_shape, requires_gpu=False):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = 3
+    kwargs['dtype'] = 'complex64'
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      num_samples = 2
+      stack_size = 3
+      length = 7
+      input_shape = (num_samples, length, stack_size)
+
+      with self.cached_session():
+        testing_utils.layer_test(
+            tfmri.layers.Conv1D,
+            kwargs=kwargs,
+            input_shape=input_shape,
+            input_dtype='complex64',
+            expected_output_shape=expected_output_shape,
+            expected_output_dtype='complex64',
+            validate_training=False)
+
+      # Training validation in `layer_test` is performed with RMSprop, with
+      # fails due to missing ops in core TF (as of TF 2.8). However, the layer
+      # can be trained with other optimizers such as SGD or Adam.
+      self._validate_training(kwargs, input_shape, 'complex64')
+
+  def _validate_training(self, layer_kwargs, input_shape, input_dtype):  # pylint: disable=missing-function-docstring
+    # Instantiate layer.
+    layer = tfmri.layers.Conv1D(**layer_kwargs)
+
+    # Create model and generate input/output data.
+    x = tf.keras.Input(shape=input_shape[1:], dtype=input_dtype)
+    model = tf.keras.Model(x, layer(x))
+    input_data = 10 * np.random.random(input_shape)
+    actual_output = model.predict(input_data)
+
+    # Now compile and train on a batch.
+    model.compile('adam', 'mse')
+    model.train_on_batch(input_data, actual_output)
+
+  def test_conv1d_complex_values(self):
+    """Test that result of complex-valued convolution is correct."""
+    input_data = tf.constant([[[ 3.+2.j],
+                               [-5.+0.j],
+                               [ 4.-3.j],
+                               [ 0.-4.j],
+                               [-5.-4.j]]], dtype=tf.complex64)  # (1, 5, 1)
+    expected_result = [[[4. - 1.j], [11. - 12.j], [0. - 5.j]]]
+
+    def kernel_initializer(shape, dtype):  # pylint: disable=unused-argument
+      return [[[-1. + 1.j]], [[0 + 0.j]], [[1. + 1.j]]]
+
+    def bias_initializer(shape, dtype):  # pylint: disable=unused-argument
+      return [2 - 3.j]
+
+    layer = tfmri.layers.Conv1D(1, 3,
+                                kernel_initializer=kernel_initializer,
+                                bias_initializer=bias_initializer,
+                                dtype='complex64')
+    layer.build((1, 5, 1))
+    result = layer(input_data)
+
+    self.assertAllClose(expected_result, result)
+
+
+@keras_parameterized.run_all_keras_modes
+class Conv2DTest(keras_parameterized.TestCase):
+
+  def _run_test(self, kwargs, expected_output_shape, spatial_shape=(7, 6)):
+    num_samples = 2
+    stack_size = 3
+    num_row, num_col = spatial_shape
+    input_data = None
+    # Generate valid input data.
+    if None in spatial_shape:
+      input_data_shape = (num_samples, num_row or 7, num_col or 6, stack_size)
+      input_data = 10 * np.random.random(input_data_shape).astype(np.float32)
+
+    with self.cached_session():
+      testing_utils.layer_test(
+          tfmri.layers.Conv2D,
+          kwargs=kwargs,
+          input_shape=(num_samples, num_row, num_col, stack_size),
+          input_data=input_data,
+          expected_output_shape=expected_output_shape)
+
+  def _run_test_extra_batch_dim(self,
+                                kwargs,
+                                expected_output_shape,
+                                spatial_shape=(7, 6)):
+    batch_shape = (2, 11)
+    stack_size = 3
+    num_row, num_col = spatial_shape
+    input_data = None
+    # Generate valid input data.
+    if None in spatial_shape:
+      input_data_shape = batch_shape + (num_row or 7, num_col or 6, stack_size)
+      input_data = 10 * np.random.random(input_data_shape).astype(np.float32)
+
+    with self.cached_session():
+      if expected_output_shape is not None:
+        expected_output_shape = (None,) + expected_output_shape
+      testing_utils.layer_test(
+          tfmri.layers.Conv2D,
+          kwargs=kwargs,
+          input_shape=batch_shape + (num_row, num_col, stack_size),
+          input_data=input_data,
+          expected_output_shape=expected_output_shape)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 5, 4, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 7, 6, 2)),
+      ('padding_same_dilation_2', {
+          'padding': 'same',
+          'dilation_rate': 2
+      }, (None, 7, 6, 2)),
+      ('strides', {
+          'strides': (2, 2)
+      }, (None, 3, 2, 2)),
+      ('dilation_rate', {
+          'dilation_rate': (2, 2)
+      }, (None, 3, 2, 2)),
+      # Only runs on GPU with CUDA, channels_first is not supported on CPU.
+      # TODO(b/62340061): Support channels_first on CPU.
+      ('data_format', {
+          'data_format': 'channels_first'
+      }, None, True),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 5, 4, 6), True),
+      ('dilation_2_unknown_width', {
+          'dilation_rate': (2, 2)
+      }, (None, None, 2, 2), False, (None, 6)),
+      ('dilation_2_unknown_height', {
+          'dilation_rate': (2, 2)
+      }, (None, 3, None, 2), False, (7, None)),
+  )
+  def test_conv2d(self,
+                  kwargs,
+                  expected_output_shape=None,
+                  requires_gpu=False,
+                  spatial_shape=(7, 6)):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = (3, 3)
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      self._run_test(kwargs, expected_output_shape, spatial_shape)
+      self._run_test_extra_batch_dim(kwargs, expected_output_shape,
+                                     spatial_shape)
+
+  def test_conv2d_regularizers(self):
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_regularizer': 'l2',
+        'bias_regularizer': 'l2',
+        'activity_regularizer': 'l2',
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv2D(**kwargs)
+      layer.build((None, 5, 5, 2))
+      self.assertEqual(len(layer.losses), 2)
+      layer(keras.backend.variable(np.ones((1, 5, 5, 2))))
+      self.assertEqual(len(layer.losses), 3)
+
+  def test_conv2d_constraints(self):
+    k_constraint = lambda x: x
+    b_constraint = lambda x: x
+
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_constraint': k_constraint,
+        'bias_constraint': b_constraint,
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv2D(**kwargs)
+      layer.build((None, 5, 5, 2))
+      self.assertEqual(layer.kernel.constraint, k_constraint)
+      self.assertEqual(layer.bias.constraint, b_constraint)
+
+  def test_conv2d_zero_kernel_size(self):
+    kwargs = {'filters': 2, 'kernel_size': 0}
+    with self.assertRaises(ValueError):
+      tfmri.layers.Conv2D(**kwargs)
+
+  def test_conv2d_invalid_output_shapes(self):
+    kwargs = {'filters': 2, 'kernel_size': 20}
+    with self.assertRaisesRegex(
+        ValueError, r"""One of the dimensions in the output is <= 0"""):
+      layer = tfmri.layers.Conv2D(**kwargs)
+      layer.build((None, 5, 5, 2))
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 5, 4, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 7, 6, 2)),
+      ('padding_same_dilation_2', {
+          'padding': 'same',
+          'dilation_rate': 2
+      }, (None, 7, 6, 2)),
+      ('strides', {
+          'strides': (2, 2)
+      }, (None, 3, 2, 2)),
+      ('dilation_rate', {
+          'dilation_rate': (2, 2)
+      }, (None, 3, 2, 2)),
+      # Only runs on GPU with CUDA, channels_first is not supported on CPU.
+      # TODO(b/62340061): Support channels_first on CPU.
+      ('data_format', {
+          'data_format': 'channels_first'
+      }, None, True),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 5, 4, 6), True),
+      ('dilation_2_unknown_width', {
+          'dilation_rate': (2, 2)
+      }, (None, None, 2, 2), False, (None, 6)),
+      ('dilation_2_unknown_height', {
+          'dilation_rate': (2, 2)
+      }, (None, 3, None, 2), False, (7, None)),
+      ('activation', {
+          'activation': 'exponential'
+      }, (None, 5, 4, 2)),
+      ('regularizer', {
+          'kernel_regularizer': 'l2'
+      }, (None, 5, 4, 2))
+  )
+  def test_conv2d_complex(self,
+                          kwargs,
+                          expected_output_shape,
+                          requires_gpu=False,
+                          spatial_shape=(7, 6)):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = 3
+    kwargs['dtype'] = 'complex64'
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      num_samples = 2
+      stack_size = 3
+      rows, cols = spatial_shape
+      input_shape = (num_samples, rows, cols, stack_size)
+
+      # Generate valid input data.
+      if None in spatial_shape:
+        input_data_shape = (num_samples, rows or 7, cols or 6, stack_size)
+        input_data = 10 * np.random.random(input_data_shape).astype(np.float32)
+      else:
+        input_data = None
+
+      with self.cached_session():
+        testing_utils.layer_test(
+            tfmri.layers.Conv2D,
+            kwargs=kwargs,
+            input_shape=input_shape,
+            input_dtype='complex64',
+            input_data=input_data,
+            expected_output_shape=expected_output_shape,
+            expected_output_dtype='complex64',
+            validate_training=False)
+
+      # Training validation in `layer_test` is performed with RMSprop, with
+      # fails due to missing ops in core TF (as of TF 2.8). However, the layer
+      # can be trained with other optimizers such as SGD or Adam.
+      self._validate_training(kwargs, input_shape, 'complex64',
+                              input_data=input_data)
+
+  def _validate_training(self, layer_kwargs, input_shape, input_dtype,  # pylint: disable=missing-function-docstring
+                         input_data=None):
+    # Instantiate layer.
+    layer = tfmri.layers.Conv2D(**layer_kwargs)
+
+    # Create model and generate input/output data.
+    x = tf.keras.Input(shape=input_shape[1:], dtype=input_dtype)
+    model = tf.keras.Model(x, layer(x))
+    if input_data is None:
+      input_data = 10 * np.random.random(input_shape)
+    actual_output = model.predict(input_data)
+
+    # Now compile and train on a batch.
+    model.compile('adam', 'mse')
+    model.train_on_batch(input_data, actual_output)
+
+
+@keras_parameterized.run_all_keras_modes
+class Conv3DTest(keras_parameterized.TestCase):
+
+  def _run_test(self, kwargs, expected_output_shape, validate_training=True):
+    num_samples = 2
+    stack_size = 3
+    num_row = 7
+    num_col = 6
+    depth = 5
+
+    with self.cached_session():
+      testing_utils.layer_test(
+          tfmri.layers.Conv3D,
+          kwargs=kwargs,
+          input_shape=(num_samples, depth, num_row, num_col, stack_size),
+          expected_output_shape=expected_output_shape,
+          validate_training=validate_training)
+
+  def _run_test_extra_batch_dim(self,
+                                kwargs,
+                                expected_output_shape,
+                                validate_training=True):
+    batch_shape = (2, 11)
+    stack_size = 3
+    num_row = 7
+    num_col = 6
+    depth = 5
+
+    with self.cached_session():
+      if expected_output_shape is not None:
+        expected_output_shape = (None,) + expected_output_shape
+
+      testing_utils.layer_test(
+          tfmri.layers.Conv3D,
+          kwargs=kwargs,
+          input_shape=batch_shape + (depth, num_row, num_col, stack_size),
+          expected_output_shape=expected_output_shape,
+          validate_training=validate_training)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 3, 5, 4, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 5, 7, 6, 2)),
+      ('strides', {
+          'strides': (2, 2, 2)
+      }, (None, 2, 3, 2, 2)),
+      ('dilation_rate', {
+          'dilation_rate': (2, 2, 2)
+      }, (None, 1, 3, 2, 2)),
+      # Only runs on GPU with CUDA, channels_first is not supported on CPU.
+      # TODO(b/62340061): Support channels_first on CPU.
+      ('data_format', {
+          'data_format': 'channels_first'
+      }, None, True),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 3, 5, 4, 6), True),
+  )
+  def test_conv3d(self, kwargs, expected_output_shape=None, requires_gpu=False):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = (3, 3, 3)
+    # train_on_batch currently fails with XLA enabled on GPUs
+    test_training = 'groups' not in kwargs or not test_util.is_xla_enabled()
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      self._run_test(kwargs, expected_output_shape, test_training)
+      self._run_test_extra_batch_dim(kwargs, expected_output_shape,
+                                     test_training)
+
+  def test_conv3d_regularizers(self):
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_regularizer': 'l2',
+        'bias_regularizer': 'l2',
+        'activity_regularizer': 'l2',
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv3D(**kwargs)
+      layer.build((None, 5, 5, 5, 2))
+      self.assertEqual(len(layer.losses), 2)
+      self.assertEqual(len(layer.losses), 2)
+      layer(keras.backend.variable(np.ones((1, 5, 5, 5, 2))))
+      self.assertEqual(len(layer.losses), 3)
+
+  def test_conv3d_constraints(self):
+    k_constraint = lambda x: x
+    b_constraint = lambda x: x
+
+    kwargs = {
+        'filters': 3,
+        'kernel_size': 3,
+        'padding': 'valid',
+        'kernel_constraint': k_constraint,
+        'bias_constraint': b_constraint,
+        'strides': 1
+    }
+    with self.cached_session():
+      layer = tfmri.layers.Conv3D(**kwargs)
+      layer.build((None, 5, 5, 5, 2))
+      self.assertEqual(layer.kernel.constraint, k_constraint)
+      self.assertEqual(layer.bias.constraint, b_constraint)
+
+  def test_conv3d_dynamic_shape(self):
+    input_data = np.random.random((1, 3, 3, 3, 3)).astype(np.float32)
+    with self.cached_session():
+      # Won't raise error here.
+      testing_utils.layer_test(
+          tfmri.layers.Conv3D,
+          kwargs={
+              'data_format': 'channels_last',
+              'filters': 3,
+              'kernel_size': 3
+          },
+          input_shape=(None, None, None, None, 3),
+          input_data=input_data)
+      if tf.test.is_gpu_available(cuda_only=True):
+        testing_utils.layer_test(
+            tfmri.layers.Conv3D,
+            kwargs={
+                'data_format': 'channels_first',
+                'filters': 3,
+                'kernel_size': 3
+            },
+            input_shape=(None, 3, None, None, None),
+            input_data=input_data)
+
+  def test_conv3d_invalid_output_shapes(self):
+    kwargs = {'filters': 2, 'kernel_size': 20}
+    with self.assertRaisesRegex(
+        ValueError, r"""One of the dimensions in the output is <= 0"""):
+      layer = tfmri.layers.Conv3D(**kwargs)
+      layer.build((None, 5, 5, 5, 2))
+
+  @parameterized.named_parameters(
+      ('padding_valid', {
+          'padding': 'valid'
+      }, (None, 3, 5, 4, 2)),
+      ('padding_same', {
+          'padding': 'same'
+      }, (None, 5, 7, 6, 2)),
+      ('strides', {
+          'strides': (2, 2, 2)
+      }, (None, 2, 3, 2, 2)),
+      ('dilation_rate', {
+          'dilation_rate': (2, 2, 2)
+      }, (None, 1, 3, 2, 2)),
+      # Only runs on GPU with CUDA, channels_first is not supported on CPU.
+      # TODO(b/62340061): Support channels_first on CPU.
+      ('data_format', {
+          'data_format': 'channels_first'
+      }, None, True),
+      # Only runs on GPU with CUDA, groups are not supported on CPU.
+      # https://github.com/tensorflow/tensorflow/issues/29005
+      ('group', {
+          'groups': 3,
+          'filters': 6
+      }, (None, 3, 5, 4, 6), True),
+      ('activation', {
+          'activation': 'exponential'
+      }, (None, 3, 5, 4, 2)),
+      ('regularizer', {
+          'kernel_regularizer': 'l2'
+      }, (None, 3, 5, 4, 2))
+  )
+  def test_conv3d_complex(self,
+                          kwargs,
+                          expected_output_shape,
+                          requires_gpu=False):
+    kwargs['filters'] = kwargs.get('filters', 2)
+    kwargs['kernel_size'] = 3
+    kwargs['dtype'] = 'complex64'
+    if not requires_gpu or tf.test.is_gpu_available(cuda_only=True):
+      num_samples = 2
+      stack_size = 3
+      rows = 7
+      cols = 6
+      depth = 5
+      input_shape = (num_samples, depth, rows, cols, stack_size)
+
+      with self.cached_session():
+        testing_utils.layer_test(
+            tfmri.layers.Conv3D,
+            kwargs=kwargs,
+            input_shape=input_shape,
+            input_dtype='complex64',
+            expected_output_shape=expected_output_shape,
+            expected_output_dtype='complex64',
+            validate_training=False)
+
+      # Training validation in `layer_test` is performed with RMSprop, with
+      # fails due to missing ops in core TF (as of TF 2.8). However, the layer
+      # can be trained with other optimizers such as SGD or Adam.
+      self._validate_training(kwargs, input_shape, 'complex64')
+
+  def _validate_training(self, layer_kwargs, input_shape, input_dtype):  # pylint: disable=missing-function-docstring
+    # Instantiate layer.
+    layer = tfmri.layers.Conv3D(**layer_kwargs)
+
+    # Create model and generate input/output data.
+    x = tf.keras.Input(shape=input_shape[1:], dtype=input_dtype)
+    model = tf.keras.Model(x, layer(x))
+    input_data = 10 * np.random.random(input_shape)
+    actual_output = model.predict(input_data)
+
+    # Now compile and train on a batch.
+    model.compile('adam', 'mse')
+    model.train_on_batch(input_data, actual_output)
+
+
+@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+class GroupedConvTest(keras_parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('Conv1D', tfmri.layers.Conv1D),
+      ('Conv2D', tfmri.layers.Conv2D),
+      ('Conv3D', tfmri.layers.Conv3D),
+  )
+  def test_group_conv_incorrect_use(self, layer):
+    with self.assertRaisesRegex(ValueError, 'The number of filters'):
+      layer(16, 3, groups=3)
+    with self.assertRaisesRegex(ValueError, 'The number of input channels'):
+      layer(16, 3, groups=4).build((32, 12, 12, 3))
+
+  @parameterized.named_parameters(
+      ('Conv1D', tfmri.layers.Conv1D, (32, 12, 32)),
+      ('Conv2D', tfmri.layers.Conv2D, (32, 12, 12, 32)),
+      ('Conv3D', tfmri.layers.Conv3D, (32, 12, 12, 12, 32)),
+  )
+  def test_group_conv(self, layer_cls, input_shape):
+    # pylint: disable=no-value-for-parameter,redundant-keyword-arg,unexpected-keyword-arg
+    if tf.test.is_gpu_available(cuda_only=True):
+      with testing_utils.use_gpu():
+        inputs = tf.random.uniform(shape=input_shape)
+
+        layer = layer_cls(16, 3, groups=4, use_bias=False)
+        layer.build(input_shape)
+
+        input_slices = tf.split(inputs, 4, axis=-1)
+        weight_slices = tf.split(layer.kernel, 4, axis=-1)
+        expected_outputs = tf.concat([
+            tf.nn.convolution(inputs, weights)
+            for inputs, weights in zip(input_slices, weight_slices)
+        ],
+                                            axis=-1)
+        self.assertAllClose(
+            layer(inputs), expected_outputs, rtol=3e-5, atol=3e-5)
+
+  def test_group_conv_depthwise(self):
+    # pylint: disable=no-value-for-parameter,redundant-keyword-arg
+    if tf.test.is_gpu_available(cuda_only=True):
+      with testing_utils.use_gpu():
+        inputs = tf.random.uniform(shape=(3, 27, 27, 32))
+
+        layer = tfmri.layers.Conv2D(32, 3, groups=32, use_bias=False)
+        layer.build((3, 27, 27, 32))
+
+        weights_dw = tf.reshape(layer.kernel, [3, 3, 32, 1])
+        expected_outputs = tf.compat.v1.nn.depthwise_conv2d(
+            inputs, weights_dw, strides=[1, 1, 1, 1], padding='VALID')
+
+        self.assertAllClose(layer(inputs), expected_outputs, rtol=1e-5)
+
+
+
+@keras_parameterized.run_all_keras_modes
+class ConvSequentialTest(keras_parameterized.TestCase):
+
+  def _run_test(self, conv_layer_cls, kwargs, input_shape1, input_shape2,
+                expected_output_shape1, expected_output_shape2):
+    kwargs['filters'] = 1
+    kwargs['kernel_size'] = 3
+    kwargs['dilation_rate'] = 2
+    with self.cached_session():
+      layer = conv_layer_cls(**kwargs)
+      output1 = layer(np.zeros(input_shape1))
+      self.assertEqual(output1.shape, expected_output_shape1)
+      output2 = layer(np.zeros(input_shape2))
+      self.assertEqual(output2.shape, expected_output_shape2)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {'padding': 'valid'},
+       (1, 8, 2), (1, 5, 2), (1, 4, 1), (1, 1, 1)),
+      ('padding_same', {'padding': 'same'},
+       (1, 8, 2), (1, 5, 2), (1, 8, 1), (1, 5, 1)),
+      ('padding_causal', {'padding': 'causal'},
+       (1, 8, 2), (1, 5, 2), (1, 8, 1), (1, 5, 1)),
+  )
+  def test_conv1d(self, kwargs, input_shape1, input_shape2,
+                  expected_output_shape1, expected_output_shape2):
+    self._run_test(tfmri.layers.Conv1D, kwargs, input_shape1, input_shape2,
+                   expected_output_shape1, expected_output_shape2)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {'padding': 'valid'},
+       (1, 7, 6, 2), (1, 6, 5, 2), (1, 3, 2, 1), (1, 2, 1, 1)),
+      ('padding_same', {'padding': 'same'},
+       (1, 7, 6, 2), (1, 6, 5, 2), (1, 7, 6, 1), (1, 6, 5, 1)),
+  )
+  def test_conv2d(self, kwargs, input_shape1, input_shape2,
+                  expected_output_shape1, expected_output_shape2):
+    self._run_test(tfmri.layers.Conv2D, kwargs, input_shape1, input_shape2,
+                   expected_output_shape1, expected_output_shape2)
+
+  @parameterized.named_parameters(
+      ('padding_valid', {'padding': 'valid'},
+       (1, 5, 7, 6, 2), (1, 8, 6, 5, 2), (1, 1, 3, 2, 1), (1, 4, 2, 1, 1)),
+      ('padding_same', {'padding': 'same'},
+       (1, 5, 7, 6, 2), (1, 8, 6, 5, 2), (1, 5, 7, 6, 1), (1, 8, 6, 5, 1)),
+  )
+  def test_conv3d(self, kwargs, input_shape1, input_shape2,
+                  expected_output_shape1, expected_output_shape2):
+    self._run_test(tfmri.layers.Conv3D, kwargs, input_shape1, input_shape2,
+                   expected_output_shape1, expected_output_shape2)
+
+  def test_dynamic_shape(self):
+    with self.cached_session():
+      layer = tfmri.layers.Conv3D(2, 3)
+      input_shape = (5, None, None, 2)
+      inputs = keras.Input(shape=input_shape)
+      x = layer(inputs)
+      # Won't raise error here with None values in input shape (b/144282043).
+      layer(x)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensorflow_mri/python/ops/convex_ops_test.py b/tensorflow_mri/python/ops/convex_ops_test.py
index 7ff22d19..9c894be7 100644
--- a/tensorflow_mri/python/ops/convex_ops_test.py
+++ b/tensorflow_mri/python/ops/convex_ops_test.py
@@ -200,7 +200,10 @@ def test_quadratic_simple(self):
     self.assertIsInstance(f.shape_tensor(), tf.Tensor)
     self.assertIsInstance(f.batch_shape_tensor(), tf.Tensor)
     self.assertAllEqual([], f.batch_shape_tensor())
-    self.assertEqual(2, f.ndim)
+    self.assertIsInstance(f.domain_dimension, int)
+    self.assertAllEqual(2, f.domain_dimension)
+    self.assertIsInstance(f.domain_dimension_tensor(), tf.Tensor)
+    self.assertAllEqual(2, f.domain_dimension_tensor())
 
   def test_quadratic_batch(self):
     """Tests `ConvexFunctionQuadratic` with batch arguments."""
@@ -222,7 +225,10 @@ def test_quadratic_batch(self):
     self.assertIsInstance(f.shape_tensor(), tf.Tensor)
     self.assertIsInstance(f.batch_shape_tensor(), tf.Tensor)
     self.assertAllEqual([2], f.batch_shape_tensor())
-    self.assertEqual(2, f.ndim)
+    self.assertIsInstance(f.domain_dimension, int)
+    self.assertAllEqual(2, f.domain_dimension)
+    self.assertIsInstance(f.domain_dimension_tensor(), tf.Tensor)
+    self.assertAllEqual(2, f.domain_dimension_tensor())
 
   def test_prox(self):
     """Tests `prox` method."""
@@ -296,21 +302,21 @@ def test_call(self):
          [4., 5., 6.]]
     x_flat = tf.reshape(x, [-1])
     reg1 = convex_ops.ConvexFunctionTotalVariation(scale=0.1,
-                                                   ndim=[2, 3],
+                                                   domain_shape=[2, 3],
                                                    axis=[0, 1])
     ref1 = 1.3
     res1 = reg1(x_flat)
     self.assertAllClose(res1, ref1)
 
     reg2 = convex_ops.ConvexFunctionTotalVariation(scale=0.1,
-                                                   ndim=[2, 3],
+                                                   domain_shape=[2, 3],
                                                    axis=1)
     res2 = reg2(x_flat)
     ref2 = 0.4
     self.assertAllClose(res2, ref2)
 
     reg3 = convex_ops.ConvexFunctionTotalVariation(scale=0.5,
-                                                   ndim=[3],
+                                                   domain_shape=[3],
                                                    axis=-1)
     res3 = reg3(x)
     ref3 = [1.0, 1.0]
diff --git a/tensorflow_mri/python/ops/optimizer_ops_test.py b/tensorflow_mri/python/ops/optimizer_ops_test.py
index 8b435901..859be9e7 100755
--- a/tensorflow_mri/python/ops/optimizer_ops_test.py
+++ b/tensorflow_mri/python/ops/optimizer_ops_test.py
@@ -97,7 +97,8 @@ def test_lasso(self):
     max_iterations = 100
 
     function_f = convex_ops.ConvexFunctionLeastSquares(operator, rhs)
-    function_g = convex_ops.ConvexFunctionL1Norm(scale=lambda_, ndim=4)
+    function_g = convex_ops.ConvexFunctionL1Norm(
+        scale=lambda_, domain_dimension=4)
 
     result = optimizer_ops.admm_minimize(function_f, function_g,
                                          absolute_tolerance=absolute_tolerance,
@@ -123,7 +124,8 @@ def test_total_variation(self):
     max_iterations = 100
 
     function_f = convex_ops.ConvexFunctionLeastSquares(operator, rhs)
-    function_g = convex_ops.ConvexFunctionL1Norm(scale=lambda_, ndim=3)
+    function_g = convex_ops.ConvexFunctionL1Norm(
+        scale=lambda_, domain_dimension=3)
 
     operator_a = linalg_ext.LinearOperatorFiniteDifference(4)
 
@@ -153,7 +155,8 @@ def test_linearized(self):
     max_iterations = 100
 
     function_f = convex_ops.ConvexFunctionLeastSquares(operator, rhs)
-    function_g = convex_ops.ConvexFunctionL1Norm(scale=lambda_, ndim=3)
+    function_g = convex_ops.ConvexFunctionL1Norm(
+        scale=lambda_, domain_dimension=3)
 
     operator_a = linalg_ext.LinearOperatorFiniteDifference(4)
 
@@ -188,7 +191,8 @@ def test_lasso_batch(self):
     max_iterations = 100
 
     function_f = convex_ops.ConvexFunctionLeastSquares(operator, rhs)
-    function_g = convex_ops.ConvexFunctionL1Norm(scale=lambda_, ndim=4)
+    function_g = convex_ops.ConvexFunctionL1Norm(
+        scale=lambda_, domain_dimension=4)
 
     result = optimizer_ops.admm_minimize(function_f, function_g,
                                          absolute_tolerance=absolute_tolerance,
diff --git a/tensorflow_mri/python/ops/recon_ops_test.py b/tensorflow_mri/python/ops/recon_ops_test.py
index 9c395646..cac99080 100755
--- a/tensorflow_mri/python/ops/recon_ops_test.py
+++ b/tensorflow_mri/python/ops/recon_ops_test.py
@@ -417,7 +417,8 @@ def test_lstsq_grasp(self, optimizer, execution_mode):  # pylint: disable=missin
 
     def _reconstruct():
       regularizer = convex_ops.ConvexFunctionTotalVariation(
-          scale=0.001, ndim=[28, 384, 384], axis=-3, dtype=tf.complex64)
+          scale=0.001, domain_shape=[28, 384, 384],
+          axis=-3, dtype=tf.complex64)
       return recon_ops.reconstruct_lstsq(
           kspace,
           image_shape=[384, 384],
diff --git a/tensorflow_mri/python/util/keras_util_test.py b/tensorflow_mri/python/util/keras_util_test.py
new file mode 100644
index 00000000..5d0da724
--- /dev/null
+++ b/tensorflow_mri/python/util/keras_util_test.py
@@ -0,0 +1,15 @@
+# Copyright 2022 University College London. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for `keras_util` module."""
diff --git a/tests/data/gre.dat b/tests/data/gre.dat
new file mode 100644
index 00000000..164bc55d
--- /dev/null
+++ b/tests/data/gre.dat
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10ce71c8cba94fb47989fe69d45c86673508f178e8b4c2399e805f27936b01ca
+size 1657344