AD with heterogeneous types in function inputs & outputs

[SemanticBeeng]

This is a a request for more information.

Currently, it is only possible to represent functions where all inputs and outputs share a single type. In future iterations, it is possible to extend support for building functions with varying input/output types and enforce constraints on both, using covariant and contravariant type bounds.

This sounds like an important limitation. 🤔

Curious: is there the same limitation in Lantern or Tangent ?

Can see reference to Nexus - very interesting.

When the shape of an N-dimensional array is known at compile-time, we can use type-level integer literals to ensure shape conforming tensor operations (inspired by Nexus and others).

"Kotlin𝛁 supports shape-safe tensor operations by encoding tensor rank as a parameter of the operand’s type signature."

It would very useful to have more details and (ideally) code examples to compare across all three frameworks and advise on future plans to advance. It would be a cool topic for a series of blog posts.

[breandan]

Hi @SemanticBeeng, thank you for your question. Let me try to clarify what is meant by,

functions where all inputs and outputs share a single type

Kotlin∇ functions are not required to have the same input and output shape, but its inputs and outputs are all of Int or all Float or all Double. I am still not sure whether implicit type conversion is a good idea, but since most programming languages allow widening or automatic type promotion (e.g. when multiplying Float by Int), we could support it.

I believe that most frameworks like TensorFlow and PyTorch default to using a single type (e.g. float32, float16) for numerical evaluation and all operations are evaluated in that context, but since mixed precision is starting to become more popular for certain applications, perhaps this is a legitimate use case. If you have a use case I would be eager to hear it.

It would very useful to have more details and (ideally) code examples to compare across all three frameworks and advise on future plans to advance. It would be a cool topic for a series of blog posts.

Agreed, it would be helpful to have side-by-side comparison of Kotlin∇/Lantern/Nexus, since they are all motivated by shape-safety/type-safety. Tangent might be out of scope, since it is a source-to-source AD for Python and has a different set of goals, but it would be great to conduct a more thorough comparison between their implementations and features.

[breandan]

The aforementioned restriction on homogenous input types has now been lifted. Previously, if you wanted to multiply a matrix variable and a scalar variable, wrapping or broadcasting was required. It is now possible to construct expressions in Kotlin∇ with heterogeneously shaped inputs. For an example of how this works, see MatrixDemo.kt, or LinearRegression.kt/MLP.kt for some practical applications.

To support this, we provide specialized datatypes which store the shape and order of the operands, e.g. products between scalars, vectors and matrices can be represented like so. Thanks to the feedback from @Redrield (cf. #4) et al. here. There is still more work to do around multivariable derivatives (e.g. I am still unhappy with the situation around vector-valued and matrix-valued gradients), but the API has gotten significantly safer and more flexible in regard to hetergeneously typed expressions.

We also provide some wrappers for type coercion where it makes sense, so it is possible to multiply primitives, Java numbers, Kotlin numbers, and other numerically typed values with functions in the Kotlin∇ DSL. Explicitly calling wrap(...) is still required in some circumstances, but most can be avoided. It may be possible to eliminate the with(Precision) pattern and infer precision using the type conversion semantics of the parent language, but this needs to be thought out more carefully.

The general philosophy of Kotlin∇ is to provide a notation that is as expressive and mathematically aligned as possible, while forbidding users from composing operations which are semantically invalid (e.g. multiplying incompatible values whose type and shape is known at compile time). We still need to publish a more thorough comparison between the available frameworks, but I wanted to keep you updated on our progress and to record some thoughts around the design process for posterity. Thanks!