Add linear algebra operations and their derivatives

[dan-zheng]

Swift for TensorFlow should provide standard linear algebra operations from tf.linalg/np.linalg:

• [x] tf.linalg.band_part
  • Done in https://github.com/tensorflow/swift-apis/pull/571.
• [x] tf.linalg.cholesky
  • Done in https://github.com/tensorflow/swift-apis/pull/563.
• [ ] tf.linalg.cholesky_solve
• [x] tf.linalg.diag (batch support: make batch from batch diagonal)
  • Done in https://github.com/tensorflow/swift-apis/pull/571.
• [x] tf.linalg.diag_part (batch support: extract batch diagonal)
  • Done in https://github.com/tensorflow/swift-apis/pull/571.
• [ ] tf.linalg.eigh
• [x] tf.linalg.logdet
  • Done in https://github.com/tensorflow/swift-apis/pull/592.
• [ ] tf.linalg.normalize
• [x] tf.linalg.qr
• [x] tf.linalg.set_diag
  • Done in https://github.com/tensorflow/swift-apis/pull/578.
• [x] tf.linalg.svd
  • Done in https://github.com/tensorflow/swift-apis/pull/578.
• [x] tf.linalg.trace
  • Done in https://github.com/tensorflow/swift-apis/pull/586.
• [x] tf.linalg.triangular_solve
  • Done in https://github.com/tensorflow/swift-apis/pull/599.
• [ ] tf.einsum
• [ ] tf.eye
  • Done in https://github.com/tensorflow/swift-apis/pull/578.

Some of these operations are already accessible via TensorFlow bindings, e.g. _Raw.cholesky, _Raw.choleskyGrad.

Reference TensorFlow derivative implementations: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_grad.py

---

Mirror of TF-980.

If anyone requests support for specific operations, please comment and we add to the checklist!

[awav]

Hello @dan-zheng, this is not a full list of operations, but it would be good to start with this (all of them implicitly should work as batched versions as well):

• [ ] Norm
• [ ] Trace
• [ ] Cholesky
• [ ] Triangular Solve
• [ ] Cholesky Solve
• [ ] Diag (make a [batch] tensor from [batch] diagonal)
• [ ] Diag Part (extract [batch] diagonal)
• [ ] Einsum
• [ ] Eye
• [ ] Logdet
• [ ] Eigh

[awav]

@dan-zheng , how can I add a new binding here? It is an auto-generated file and I don't see how it was generated.

[Shashi456]

@awav you can't add new bindings in that file, You can check where how its generated here.

[awav]

[EDITED] that's my bad, "TriangularSolve" is in the RawOpsGenerated.swift file, as well as many other linalg ops.

@Shashi456 , yeah, I've checked it. But, it is not clear why for e.g. "TriangularSolve" is not reflected in the RawOpsGenerated.swift file, but "TriangularSolve" is in the ops.pbtxt. Does S4TF uses different protobuf file?

[dan-zheng]

@Shashi456 , yeah, I've checked it. But, it is not clear why for e.g. "TriangularSolve" is not reflected in the RawOpsGenerated.swift file, but "TriangularSolve" is in the ops.pbtxt. Does S4TF uses different protobuf file?

Swift for TensorFlow uses the same ops.pbtxt - here's the script that generates op bindings.

TensorFlow.Raw and ops.pbtxt are consistent, both define matrixTriangularSolve and batchMatrixTriangularSolve:

[Shashi456]

@awav are you working on supporting these ops? could i ask what ops you are working on currently?

[awav]

@Shashi456, solvers and trace for now

@dan-zheng, you can mark diag and diag_part as finished. They were implemented in https://github.com/tensorflow/swift-apis/pull/571

[awav]

@Shashi456, @dan-zheng hello guys, sorry for the delay with PRs for solvers and other promised ops. I'm at NeurIPS and don't have much time to finish the work. Thanks!

[awav]

@dan-zheng, trace can be marked as done now.

[pedronahum]

Hi, If I can add to the wishlist, the following matrix operations could be useful:

• QR eigen decomposition with explicit shift (Wilkinson shif)
• the Schur decomposition