Repository: google/jax

Issue:
scipy.linalg.schur decomposition
Looks like schur decomposition isn't implemented at the moment. Are there any plans to support this? Also, anything specific that might block me from implementing this?


Discussion:
This sounds pretty reasonable to me. I don't think there are any current plan to implement this but we would welcome contributions.

See these files for examples of how we wrap LAPACK/CuSolver for new low-level solvers:
https://github.com/google/jax/blob/master/jaxlib/lapack.pyx
https://github.com/google/jax/blob/master/jaxlib/cusolver.py

(I'm guessing we don't want to implement this method ourselves)
Implemented a wrapper for the required lapack routine, but I'm facing several issues while building jaxlib from source to incorporate these changes. Are jaxlib wheels generated by the CI/CD if I open a PR? That way I can avoid generating them locally.

Tested in ubuntu 18.04 (docker).
No, the github CI uses the `_minimum_jaxlib_version` specified in `setup.py` and does not cover changes to jaxlib sources. There is some testing involving jaxlib compilation outside github CI once the `Pull Ready` tag is added, but that is unfortunately fairly opaque to non-Google employees. The best course would be to figure out how to test your changes locally.
@ayush-1506 But we can probably help you figure out what's going wrong with your jaxlib build. Have you seen https://jax.readthedocs.io/en/latest/developer.html#building-from-source ?
@hawkinsp Yes, I'm following the same instructions. The error seems trivial:
```
Server terminated abruptly (error code: 14, error message: 'Socket closed', log file: '/root/.cache/bazel/_bazel_root/deb80d6610824a92deeac7b7fd0f3e3c/server/jvm.out')
```

My guess is that this is happening due to oom. 
@ayush-1506 Yes, that is plausible. You might try setting the `--jobs` argument to Bazel; using fewer jobs than you have cores might help if you are RAM constrained. You can pass extra Bazel options through the `build.py` script using `--bazel_options="--jobs 4"` or something like that.
@hawkinsp Thanks, reducing the number of jobs seems to work. However, I ran into another issue after this. (This doesn't seem to be related to oom). I'm building this without rocm/cuda/tpu.

```
[0 / 17] [Prepa] BazelWorkspaceStatusAction stable-status.txt
[136 / 1,670] Compiling com_google_protobuf/src/google/protobuf/compiler/java/java_message_lite.cc; 1s local ... (2 actions, 1 running)
[348 / 1,871] Compiling com_google_absl/absl/time/internal/cctz/src/time_zone_info.cc; 1s local ... (2 actions, 1 running)
[534 / 1,871] Compiling org_tensorflow/tensorflow/core/platform/hash.cc; 1s local ... (2 actions, 1 running)
[805 / 1,871] Compiling external/org_tensorflow/tensorflow/core/framework/full_type.pb.cc; 0s local ... (2 actions, 1 running)
[933 / 1,871] Compiling org_tensorflow/tensorflow/core/platform/default/posix_file_system.cc; 2s local ... (2 actions, 1 running)
ERROR: /root/.cache/bazel/_bazel_root/deb80d6610824a92deeac7b7fd0f3e3c/external/org_tensorflow/tensorflow/core/BUILD:1593:16: C++ compilation of rule '@org_tensorflow//tensorflow/core:framework_internal_impl' failed (Exit 1): gcc failed: error executing command
  (cd /root/.cache/bazel/_bazel_root/deb80d6610824a92deeac7b7fd0f3e3c/execroot/__main__ && \
  exec env - \
    PATH=/usr/local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \
    PWD=/proc/self/cwd \
    TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 \
    TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 \
  /usr/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections '-std=c++11' -MD -MF bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/core/_objs/framework_internal_impl/batch_util.pic.d '-frandom-seed=bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/core/_objs/framework_internal_impl/batch_util.pic.o' -fPIC -DHAVE_SYS_UIO_H -DTF_USE_SNAPPY -DEIGEN_MPL2_ONLY '-DEIGEN_MAX_ALIGN_BYTES=64' -iquoteexternal/org_tensorflow -iquotebazel-out/k8-opt/bin/external/org_tensorflow -iquoteexternal/com_google_protobuf -iquotebazel-out/k8-opt/bin/external/com_google_protobuf -iquoteexternal/zlib -iquotebazel-out/k8-opt/bin/external/zlib -iquoteexternal/eigen_archive -iquotebazel-out/k8-opt/bin/external/eigen_archive -iquoteexternal/com_google_absl -iquotebazel-out/k8-opt/bin/external/com_google_absl -iquoteexternal/nsync -iquotebazel-out/k8-opt/bin/external/nsync -iquoteexternal/gif -iquotebazel-out/k8-opt/bin/external/gif -iquoteexternal/libjpeg_turbo -iquotebazel-out/k8-opt/bin/external/libjpeg_turbo -iquoteexternal/com_googlesource_code_re2 -iquotebazel-out/k8-opt/bin/external/com_googlesource_code_re2 -iquoteexternal/farmhash_archive -iquotebazel-out/k8-opt/bin/external/farmhash_archive -iquoteexternal/fft2d -iquotebazel-out/k8-opt/bin/external/fft2d -iquoteexternal/highwayhash -iquotebazel-out/k8-opt/bin/external/highwayhash -iquoteexternal/double_conversion -iquotebazel-out/k8-opt/bin/external/double_conversion -iquoteexternal/snappy -iquotebazel-out/k8-opt/bin/external/snappy -isystem external/com_google_protobuf/src -isystem bazel-out/k8-opt/bin/external/com_google_protobuf/src -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib -isystem external/org_tensorflow/third_party/eigen3/mkl_include -isystem bazel-out/k8-opt/bin/external/org_tensorflow/third_party/eigen3/mkl_include -isystem external/eigen_archive -isystem bazel-out/k8-opt/bin/external/eigen_archive -isystem external/nsync/public -isystem bazel-out/k8-opt/bin/external/nsync/public -isystem external/gif -isystem bazel-out/k8-opt/bin/external/gif -isystem external/farmhash_archive/src -isystem bazel-out/k8-opt/bin/external/farmhash_archive/src -isystem external/double_conversion -isystem bazel-out/k8-opt/bin/external/double_conversion -Wno-sign-compare -Wno-stringop-truncation -mavx '-std=c++14' -DEIGEN_AVOID_STL_ARRAY -Iexternal/gemmlowp -Wno-sign-compare '-ftemplate-depth=900' -fno-exceptions -DINTEL_MKL -msse3 -DTENSORFLOW_MONOLITHIC_BUILD -pthread '-DINTEL_MKL=1' -fno-canonical-system-headers -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/org_tensorflow/tensorflow/core/util/batch_util.cc -o bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/core/_objs/framework_internal_impl/batch_util.pic.o)
Execution platform: @local_execution_config_platform//:platform
In file included from external/org_tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/FixedPoint:46,
                 from external/org_tensorflow/tensorflow/core/framework/numeric_types.h:24,
                 from external/org_tensorflow/tensorflow/core/framework/allocator.h:26,
                 from external/org_tensorflow/tensorflow/core/framework/tensor.h:23,
                 from external/org_tensorflow/tensorflow/core/util/batch_util.h:18,
                 from external/org_tensorflow/tensorflow/core/util/batch_util.cc:16:
external/org_tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/PacketMathAVX.h:14:41: warning: ignoring attributes on template argument '__m256i' {aka '__vector(4) long long int'} [-Wignored-attributes]
 typedef eigen_packet_wrapper<__m256i, 10> Packet32q8i;
                                         ^
external/org_tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/PacketMathAVX.h:15:41: warning: ignoring attributes on template argument '__m128i' {aka '__vector(2) long long int'} [-Wignored-attributes]
 typedef eigen_packet_wrapper<__m128i, 11> Packet16q8i;
                                         ^
gcc: fatal error: Killed signal terminated program cc1plus
compilation terminated.
Target //build:build_wheel failed to build
INFO: Elapsed time: 1050.784s, Critical Path: 92.65s
INFO: 974 processes: 43 internal, 931 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
Traceback (most recent call last):
  File "build/build.py", line 521, in <module>
    main()
  File "build/build.py", line 516, in main
    shell(command)
  File "build/build.py", line 51, in shell
    output = subprocess.check_output(cmd)
  File "/usr/local/lib/python3.8/subprocess.py", line 411, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/usr/local/lib/python3.8/subprocess.py", line 512, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-3.7.2-linux-x86_64', 'run', '--verbose_failures=true', '--jobs=2', '--config=short_logs', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/root/jax/dist']' returned non-zero exit status 1.
```
Wait, I think I figured it out. Generated the wheels successfully. However `pip3 install dist/jaxlib-0.1.66-cp36-none-manylinux2010_x86_64.whl` now gives me `jaxlib-0.1.66-cp36-none-manylinux2010_x86_64.whl is not a supported wheel on this platform.`. 
I'm on x86-64 and python3.6.9

Why could this be happening?
@hawkinsp 
Ok, never mind that. Everything works now.
I'm nearly done with the lapack wrapper (was away for a couple of weeks). Is `CustomCallWithLayout` documented somewhere? Can't find any docs for it. I'm trying to understand what its arguments mean. Can't wrap my head around what shape `shape_with_layout` refers to.

@jakevdp @hawkinsp 

I don't think there is any documentation for `CustomCallWithLayout`; your best bet is probably to take a look at examples within the JAX codebase.

As for `shape_with_layout`, this specifies the shape of the outputs. It's essentially a set of abstract shapes that specify the dtype, the dimensions, and the memory layout. e.g. for a 2D array, row-major layout is `(0, 1)` and column-major layout is `(1, 0)`.

Hope that helps
Hi everyone,

I landed on this issue trying to implement the `logm` matrix function for https://github.com/google/jax/issues/5469.

I didn't know if it was still active so went ahead wrote a wrapper for LAPACK's gees routine (which computes the Schur decomposition). It uses C++ templates like the refactored lapack wrappers in `jaxlib/lapack_kernels.cc`. 

With these additions, on my macOS machine, jaxlib compiles. I wrote a schur primitive (by closely following the eig primitive code) to use the wrapper.

On preliminary tests it outputs the same decomposition as scipy's schur function, so it appears to work.  

I am currently working on writing more tests for this decomposition. The current version of my code passes the existing linalg tests so it looks like nothing has broken. I thought I would write before writing the test suite to maybe get some early feedback.

The signatures for the wrappers are [here](https://github.com/SaturdayGenfo/jax/blob/3c17096bc2dfa4aa9c04028f61f58f3b2b8e8df5/jaxlib/lapack_kernels.h#L152). (For reference, the LAPACK I signatures I followed are [here](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-fortran/top/lapack-routines/lapack-least-squares-and-eigenvalue-problem-routines/lapack-least-squares-and-eigenvalue-problem-driver-routines/nonsymmetric-eigenvalue-problems-lapack-driver-routines/gees.html))
The implementation of the wrappers is [here](https://github.com/SaturdayGenfo/jax/blob/dfd69b3cc20000058366f1d8858c4bfcbae17f8a/jaxlib/lapack_kernels.cc#L631).
The tie-in to scipy's LAPACK is [here](https://github.com/SaturdayGenfo/jax/blob/3c17096bc2dfa4aa9c04028f61f58f3b2b8e8df5/jaxlib/lapack.cc#L118).
The python gees function is [here](https://github.com/SaturdayGenfo/jax/blob/dfd69b3cc20000058366f1d8858c4bfcbae17f8a/jaxlib/lapack.py#L588)
The primitive is [here](https://github.com/SaturdayGenfo/jax/blob/dfd69b3cc20000058366f1d8858c4bfcbae17f8a/jax/_src/lax/linalg.py#L1446).

I used function pointers for the *optional* callable argument that `gees` uses.  Is there a way to pass a callable through `CustomCallWithLayout` ? Since I didn't know how to refer to a callable in the `operands` argument, using the callable argument raises an error for now.






Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -42,6 +42,7 @@
 
 from jax._src.lib import xla_client
 from jax._src.lib import xla_bridge as xb
+from jax._src.lib import version as jaxlib_version
 
 xops = xla_client.ops
 
@@ -1441,3 +1442,119 @@ def tridiagonal_solve(dl, d, du, b):
     raise ValueError(f'Only f32/f64 are supported, got {t}')
 
   return tridiagonal_solve_p.bind(dl, d, du, b, m=m, n=n, ldb=ldb, t=t)
+
+
+# Schur Decomposition
+
+
+def schur(x,
+          compute_schur_vectors=True,
+          sort_eig_vals=False,
+          select_callable=None):
+  return schur_p.bind(
+      x,
+      compute_schur_vectors=compute_schur_vectors,
+      sort_eig_vals=sort_eig_vals,
+      select_callable=select_callable)
+
+
+def _schur_impl(operand, *, compute_schur_vectors, sort_eig_vals,
+                select_callable):
+  return xla.apply_primitive(
+      schur_p,
+      operand,
+      compute_schur_vectors=compute_schur_vectors,
+      sort_eig_vals=sort_eig_vals,
+      select_callable=select_callable)
+
+
+def _schur_translation_rule(c, operand, *, compute_schur_vectors,
+                            sort_eig_vals):
+  raise NotImplementedError(
+      "Schur decomposition is only implemented on the CPU backend.")
+
+
+def _schur_abstract_eval(operand, *, compute_schur_vectors, sort_eig_vals,
+                         select_callable):
+
+  if operand.ndim < 2 or operand.shape[-2] != operand.shape[-1]:
+    raise ValueError("Argument to Schur decomposition must have "
+                     "shape [..., n, n], got shape {}".format(operand.shape))
+
+  batch_dims = operand.shape[:-2]
+  n = operand.shape[-1]
+  dtype = operand.dtype
+  dtype = dtypes.canonicalize_dtype(dtype)
+  T = operand.update(shape=batch_dims + (n, n), dtype=dtype)
+  vs = operand.update(shape=batch_dims + (n, n), dtype=dtype)
+
+  return (T, vs) if compute_schur_vectors else (T,)
+
+
+def _schur_cpu_translation_rule(c, operand, *, compute_schur_vectors,
+                                sort_eig_vals, select_callable):
+  shape = c.get_shape(operand)
+  batch_dims = shape.dimensions()[:-2]
+
+  if jaxlib_version < (0, 1, 72):
+    raise NotImplementedError(
+        "The Schur primitive is only implemented for jaxlib versions >= 0.1.72"
+    )
+
+  _cpu_gees = lapack.gees
+
+  if sort_eig_vals:
+    T, vs, sdim, info = _cpu_gees(
+        c,
+        operand,
+        jobvs=compute_schur_vectors,
+        sort=sort_eig_vals,
+        select=select_callable)
+  else:
+    T, vs, info = _cpu_gees(
+        c,
+        operand,
+        jobvs=compute_schur_vectors,
+        sort=sort_eig_vals,
+        select=select_callable)
+
+  ok = xops.Eq(info, xops.ConstantLiteral(c, np.array(0, np.int32)))
+  T = _broadcasting_select(c, xops.Reshape(ok, batch_dims + (1, 1)), T,
+                           _nan_like(c, T))
+  output = [T]
+  if compute_schur_vectors:
+    vs = _broadcasting_select(c, xops.Reshape(ok, batch_dims + (1, 1)), vs,
+                              _nan_like(c, vs))
+
+    output.append(vs)
+
+  return xops.Tuple(c, output)
+
+
+def _schur_batching_rule(batched_args, batch_dims, *, compute_schur_vectors,
+                         sort_eig_vals, select_callable):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.moveaxis(x, bd, 0)
+
+  return schur_p.bind(
+      x,
+      compute_schur_vectors=compute_schur_vectors,
+      sort_eig_vals=sort_eig_vals,
+      select_callable=select_callable), (0,) * (1 + compute_schur_vectors)
+
+
+def _schur_jvp_rule(primals, tangents, *, compute_schur_vectors, sort_eig_vals):
+  raise NotImplementedError(
+      'The differentiation rules for the Schur factorization have not been implemented.'
+  )
+
+
+schur_p = Primitive('schur')
+schur_p.multiple_results = True
+schur_p.def_impl(_schur_impl)
+schur_p.def_abstract_eval(_schur_abstract_eval)
+xla.translations[schur_p] = _schur_translation_rule
+xla.backend_specific_translations['cpu'][schur_p] = _schur_cpu_translation_rule
+batching.primitive_batchers[schur_p] = _schur_batching_rule
+ad.primitive_jvps[schur_p] = _schur_jvp_rule
diff --git a/jax/lax/linalg.py b/jax/lax/linalg.py
--- a/jax/lax/linalg.py
+++ b/jax/lax/linalg.py
@@ -31,4 +31,6 @@
   triangular_solve_p,
   tridiagonal_solve,
   tridiagonal_solve_p,
+  schur,
+  schur_p
 )
diff --git a/jaxlib/lapack.py b/jaxlib/lapack.py
--- a/jaxlib/lapack.py
+++ b/jaxlib/lapack.py
@@ -584,3 +584,79 @@ def geev(c, a, jobvl=True, jobvr=True):
   else:
     return (_ops.GetTupleElement(out, 2), _ops.GetTupleElement(out, 3),
             _ops.GetTupleElement(out, 4), _ops.GetTupleElement(out, 5))
+
+# # gees : Schur factorization
+
+def gees(c, a, jobvs=True, sort=False, select=None):
+  a_shape = c.get_shape(a)
+  dtype = a_shape.element_type()
+  dims = a_shape.dimensions()
+  assert len(dims) >= 2
+  m, n = dims[-2:]
+  assert m == n
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  b = 1
+  for d in batch_dims:
+    b *= d
+  layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
+
+  if sort:
+    raise NotImplementedError(
+        "The sort feature of LAPACK's gees routine is not implemented.")
+
+  jobvs = ord('V' if jobvs else 'N')
+  sort = ord('S' if sort else 'N')
+
+  if dtype == np.float32 or dtype == np.float64:
+    fn = b"lapack_sgees" if dtype == np.float32 else b"lapack_dgees"
+    schurvecs_type = dtype
+    workspaces = (Shape.array_shape(np.dtype(schurvecs_type), dims, layout),)
+    eigvals = (Shape.array_shape(
+        np.dtype(dtype), batch_dims + (n,), tuple(range(num_bd, -1, -1))),
+               Shape.array_shape(
+                   np.dtype(dtype), batch_dims + (n,),
+                   tuple(range(num_bd, -1, -1))))
+  elif dtype == np.complex64 or dtype == np.complex128:
+    fn = b"lapack_cgees" if dtype == np.complex64 else b"lapack_zgees"
+    schurvecs_type = dtype
+    workspaces = (
+        Shape.array_shape(np.dtype(schurvecs_type), dims, layout),
+        Shape.array_shape(
+            np.dtype(np.float32 if dtype == np.complex64 else np.float64),
+            (n,), (0,)))
+    eigvals = (Shape.array_shape(
+        np.dtype(dtype), batch_dims + (n,), tuple(range(num_bd, -1, -1))),)
+  else:
+    raise NotImplementedError("Unsupported dtype {}".format(dtype))
+
+  out = _ops.CustomCallWithLayout(
+      c,
+      fn,
+      operands=(
+          _constant_s32_scalar(c, b),
+          _constant_s32_scalar(c, n),
+          _ops.Constant(c, np.uint8(jobvs)),
+          _ops.Constant(c, np.uint8(sort)),
+          #figure out how to put the callable select function here
+          a),
+      shape_with_layout=Shape.tuple_shape(workspaces + eigvals + (
+          Shape.array_shape(np.dtype(schurvecs_type), dims, layout),
+          Shape.array_shape(
+              np.dtype(np.int32), batch_dims, tuple(range(num_bd - 1, -1, -1))),
+          Shape.array_shape(
+              np.dtype(np.int32), batch_dims, tuple(range(num_bd -
+                                                          1, -1, -1))))),
+      operand_shapes_with_layout=(
+          Shape.array_shape(np.dtype(np.int32), (), ()),
+          Shape.array_shape(np.dtype(np.int32), (), ()),
+          Shape.array_shape(np.dtype(np.uint8), (), ()),
+          Shape.array_shape(np.dtype(np.uint8), (), ()),
+          Shape.array_shape(dtype, dims, layout),
+      ))
+  if sort == ord('S'):
+    return (_ops.GetTupleElement(out, 0), _ops.GetTupleElement(out, 3),
+            _ops.GetTupleElement(out, 4), _ops.GetTupleElement(out, 5))
+  else:
+    return (_ops.GetTupleElement(out, 0), _ops.GetTupleElement(out, 3),
+            _ops.GetTupleElement(out, 5))


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -1440,5 +1440,42 @@ def test_tridiagonal_solve(self, dtype):
     A[[0, 1], [1, 2]] = du[:-1]
     np.testing.assert_allclose(A @ X, B, rtol=1e-6, atol=1e-6)
 
+  @parameterized.named_parameters(
+        jtu.cases_from_list({
+            "testcase_name":
+            "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
+            "shape": shape, "dtype": dtype
+        } for shape in [(4, 4), (15, 15), (50, 50), (100, 100)]
+                            for dtype in float_types + complex_types))
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testSchur(self, shape, dtype):
+      if jax._src.lib.version < (0, 1, 72):
+          self.skipTest("Schur LAPACK wrapper only implemented for jaxlib versions >= 0.1.72")
+      rng = jtu.rand_default(self.rng())
+      args_maker = lambda: [rng(shape, dtype)]
+
+      self._CheckAgainstNumpy(osp.linalg.schur, lax.linalg.schur, args_maker)
+      self._CompileAndCheck(lax.linalg.schur, args_maker)
+
+  @parameterized.named_parameters(
+      jtu.cases_from_list({
+          "testcase_name":
+          "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
+          "shape": shape, "dtype": dtype
+      } for shape in [(2, 2), (4, 4), (15, 15), (50, 50), (100, 100)]
+                          for dtype in float_types + complex_types))
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testSchurBatching(self, shape, dtype):
+      if jax._src.lib.version < (0, 1, 72):
+          self.skipTest("Schur LAPACK wrapper only implemented for jaxlib versions >= 0.1.72")
+      rng = jtu.rand_default(self.rng())
+      batch_size = 10
+      shape = (batch_size, ) + shape
+      args = rng(shape, dtype)
+      reconstruct = vmap(lambda S, T: S @ T @ jnp.conj(S.T))
+
+      Ts, Ss = vmap(lax.linalg.schur)(args)
+      self.assertAllClose(reconstruct(Ss, Ts), args, atol=1e-4)
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())