Function builder cannot infer generic parameters even though direct call to `buildBlock` can

Using Swift
,
1

Hi folks,

I got a question about a function builder diagnostic.

Our deep learning library defines a Sequential type for sequentially composing neural network layers. It's powered by a function builder called LayerBuilder.

Sequential is used like this:

let model = Sequential {
    Conv2D<Float>(...) // first layer, feeding output to:
    AvgPool2D<Float>(...) // second layer, feeding output to...
    Flatten<Float>()
    Dense<Float>(...)
    Dense<Float>(...)
}

And ideally, we'd like the function builder to infer layer generic parameters so that only one needs to be specified:

let model = Sequential {
    Conv2D<Float>(...) // only one `<Float>` specified
    AvgPool2D(...)
    Flatten()
    Dense(...)
    Dense(...)
}

But it doesn't work. Here's the error:

seq.swift:43:27: error: static method 'buildBlock' requires the types 'Dense<Float>.Output' (aka 'Tensor<Float>') and 'Dense<_>.Input' (aka 'Tensor<_>') be equivalent
let modelBad = Sequential {
                          ^
seq.swift:30:15: note: where 'L1.Output' = 'Dense<Float>.Output' (aka 'Tensor<Float>'), 'L2.Input' = 'Dense<_>.Input' (aka 'Tensor<_>')
  static func buildBlock<L1: Layer, L2: Layer>(_ l1: L1, _ l2: L2) -> Sequential<L1, L2>
              ^
seq.swift:45:3: error: generic parameter 'Scalar' could not be inferred
  Dense()
  ^
seq.swift:7:14: note: 'Scalar' declared as parameter to type 'Dense'
struct Dense<Scalar>: Layer {
             ^
seq.swift:45:3: note: explicitly specify the generic arguments to fix this issue
  Dense()
  ^
       <Any>

Is this expected behavior?

Direct calls to LayerBuilder.buildBlock do type-check when some generic parameters are inferred, which made me wonder whether this is a type inference deficiency specific to function builders:

// A direct call to the function builder entry point does compile:
_ = LayerBuilder.buildBlock(
  // generic parameter specified only once:
  Dense<Float>(),
  Dense())

Minimized reproducer

Uncomment FIXME code for error.

protocol Layer {
  associatedtype Input
  associatedtype Output
}

struct Tensor<Scalar> {}
struct Dense<Scalar>: Layer {
  typealias Input = Tensor<Scalar>
  typealias Output = Tensor<Scalar>
}

struct Sequential<L1: Layer, L2: Layer>: Layer {
  var layer1: L1, layer2: L2

  typealias Input = L1.Input
  typealias Output = L2.Output

  init(_ layer1: L1, _ layer2: L2) {
    self.layer1 = layer1
    self.layer2 = layer2
  }

  init(@LayerBuilder layers: () -> Self) {
    self = layers()
  }
}

@_functionBuilder
struct LayerBuilder {
  static func buildBlock<L1: Layer, L2: Layer>(_ l1: L1, _ l2: L2) -> Sequential<L1, L2>
  where L1.Output == L2.Input {
    Sequential(l1, l2)
  }
}

// Example:
let model = Sequential {
  Dense<Float>()
  Dense<Float>()
}

// FIXME: this doesn't compile:
// let modelBad = Sequential {
//   Dense<Float>()
//   Dense()
// }

// Ideally, we'd like to infer generic parameters so that they only need to be specified once.
//
// Not ideal:
//
//     let model = Sequential {
//       Conv2D<Float>(...)
//       AvgPool2D<Float>(...)
//       Flatten<Float>()
//       Dense<Float>(...)
//       Dense<Float>(...)
//     }
//
// Ideal:
//
//     let model = Sequential {
//       Conv2D<Float>(...) // generic parameter specified only once
//       AvgPool2D(...)
//       Flatten()
//       Dense(...)
//       Dense(...)
//     }

// A direct call to the function builder entry point does compile:
_ = LayerBuilder.buildBlock(
  // generic parameter specified only once:
  Dense<Float>(),
  Dense())

2

I'm happy to file an issue at https://bugs.swift.org!

Wanted to share here first so folks can help sanity check, in case the error is expected behavior.

3

Fwiw, I tried to do something similar and also found that it doesn’t work. It would be great if this limitation could be lifted eventually.

1 Like
4

Glad to hear I'm not alone in encountering this!

I cross-posted this issue to [SR-12682] Function builder generic parameter type inference fails · Issue #55126 · apple/swift · GitHub. Hopefully type checker experts can comment sometime.

5

Yeah, this is expected behavior. The type inferred for an expression in a function builder does not affect the types inferred for the other expressions.

Unlike most expressions in Swift, function builders use something called one-way constraints for type inference (this is why a regular function call to buildBlock behaves differently wrt type inference). A one-way constraint means that type information can flow only in one direction, instead of the usual bi-directional type inference. So, in this case, the type information Conv2D<Float> will affect the overall type inferred for model, but not for the other expressions inside Sequential { ....

@Douglas_Gregor wrote up an explanation of one-way constraints and why they're used in his PR that added them: [Constraint solver] Introduce one-way binding constraints. by DougGregor · Pull Request #25983 · apple/swift · GitHub.

2 Likes
6

Thanks for the explanation Holly!

"What are one-way constraints" makes sense, but the motivation behind them isn't clear to me. Could you please clarify? Are they important for (1) correct function builder type inference or (2) is it a performance optimization?

The linked PR shares an example of how one-way constraints prune the constraint solver solution space, but I can't tell if that pruning falls under (1) or (2).

7

The motivation is both to 1) better match the function builder mental model, and 2) to improve performance by allowing the constraint system for a function builder to be split up into smaller components that can be solved independently.

  1. The basic idea of a function builder is that certain functions implicitly build up a value using the results of the function's statements. Using one-way constraints for each of the function's statements better matches how regular functions bodies are type-checked across statements, and it matches the mental model of aggregating a sequence of independent components.

  2. One-way constraints are a huge performance win because each component within the function builder body can be type-checked separately.

3 Likes
8

In the use case I have, I ended up abandoning function builders and using a chained builder-style API instead. This approach inferred well in Swift 5.1 although I’m having some trouble with it in Swift 5.2 (still trying to get to sort this out).

If inference is able to work with the chained syntax then omitting it from function builders doesn’t prevent a library design from expressing something equivalent. It just removes function builder syntax as an option for the library. IMO, this is an unfortunate tradeoff to run into.

3 Likes
9

What Holly says is exactly right. Arbitrary type information not back-propagating and side-propagating between statements is an intentional aspect of the function builder design. One of the reasons we didn't put function builders through evolution last fall was specifically that the implementation could have propagated types that way, which we didn't consider acceptable or sustainable for the feature to commit to.

Part of our thinking here is that we've been conscious of the possibility of type-checking all closures in their surrounding context the same way that function-builder closures are. That would eliminate the problem where multi-statement closures passed to generic functions often need to declare their signature. But to consider that, we need to (1) make sure that the type-checking model isn't vastly different from what it's been before and (2) stave off as many sources of exponential complexity as we can. That means restricting how types propagate between statements to be as consistent as possible with a statement-by-statement model.

4 Likes
10

Good idea! A fluent builder-style API seems like a viable alternative that should support type inference in the original example.

I agree with this sentiment. I'd like to use function builders with the desired type inference, but it sounds like function builders are intentionally designed to not support it (for well-considered reasons).

11

I see - I've heard of this direction to generalize one-way constraints to all closures, thanks for all the clarification.

Regarding multi-statement closures, the tradeoff seems to be:

  1. Explicit closure type annotation, but maximum type inference within the body. This is the status quo.
  2. Explicit type annotation for variable bindings within the body, but inferred closure type. This would be enabled by one-way constraints.

It's hard for me personally to judge whether (1) or (2) is better.

(1) fits the principle of "no inference for function signatures (where annotation matters, because we want users to be explicit to prevent unexpected behavior) but maximum inference inside the body (where annotation matters much less)".

But (2) seems to enable type inference to work for some code that is currently diagnosed with the following, which I run into somewhat frequently:

unable to infer closure return type; add explicit type to disambiguate
12

Understood. What isn't clear to me though is why this is necessary in function builders if an analogous fluent builder API doesn't run into the same issues. Is there a non-obvious difference in type checking consequences between the two approaches to syntax? Or is the intent to discourage libraries that would rely on this kind of inference altogether?

1 Like
13

I think these are good nuanced questions.

Also: I think the "fluent builder API" is just explicit calls to the function builder's buildBlock. For my use case, the function builder is syntactic sugar for a fluent builder API.

1 Like
14

You wouldn't have to add explicit type annotations within the body; that would clearly not be consistent with current type-checker behavior.

A one-way constraint is like the status quo for type-checking a variable binding:

let x = 15
return x + UInt(27)

This doesn't type-check because type information only flows "one way" through the variable binding: we decide that x: Int without considering that x is required to be convertible to UInt on the next line.

Another way of understanding the use of one-way constraints in function builders is that they're as if the closure was transformed to assign all the statement results into locals. So given this closure:

Sequential { 
  Conv2D<Float>(...)
  AvgPool2D<Float>(...)
  Flatten<Float>()
  Dense<Float>(...)
  Dense<Float>(...)
}

Instead of thinking of the transform like this:

Sequential {
  LayerBuilder.buildBlock(
    Conv2D<Float>(...),
    AvgPool2D<Float>(...),
    Flatten<Float>(),
    Dense<Float>(...),
    Dense<Float>(...))
}

Think of it more like this:

Sequential { 
  let a = Conv2D<Float>(...)
  let b = AvgPool2D<Float>(...)
  let c = Flatten<Float>()
  let d = Dense<Float>(...)
  let e = Dense<Float>(...)
  return LayerBuilder.buildBlock(a, b, c, d, e)
}
13 Likes
15

(moving bugs.swift.org discussion here, upon request)

Thanks for the technical explanations John! The "lack of type inference" is expected behavior given how function builder syntax is expanded into variable bindings.

I'm still curious about answers to @anandabits's use-case-motivated questions from above:

Should libraries give up on ever using function builders for this style of type inference? Are libraries encouraged to use explicit generic parameters or a different pattern (e.g. fluent method-chaining builders) instead?

// Will this `<Float>` generic parameter type inference never be possible?
let model = Sequential {
    Conv2D<Float>(...) // first layer, feeding output to:
    AvgPool2D(...) // AvgPool2D<Float>, feeding output to...
    Flatten() // Flatten<Float>
    Dense(...) // Dense<Float>
    Dense(...) // Dense<Float>
}
16

You might be able to get this kind of propagation if the builder type itself is parameterized and you use buildExpression, which intentionally does fire “before the binding”. I believe we’d consider that intended behavior.

17

That doesn’t introduce a way to flow type information from one builder expression to the next, does it? That is the capability that is present in fluent-style syntax.

My original question was whether the fluent-style API design somehow reduces work for the type checker in a way that function builders would not, or whether the same issues with inference and type checking performance would exist with fluent-style syntax. I’m still curious to know the answer, especially if there are potential pitfalls we should be aware of when designing fluent-style API which do flow type information forward.

18

It would allow type information to flow between the builder type and the element expressions, which effectively means between all element expressions simultaneously. It would not be sufficient for writing e.g. a functional/monadic chain where the output type of one element interacts with the input type of the next.

1 Like
19

By allowing that kind of typing interaction, however, it might significantly impact the performance of type-checking such a builder. I’d want @Douglas_Gregor’s opinion about whether we consider it intended.

20

Gotcha, I was wondering about that.

Ahh, gotcha. I'm not thinking of a monad specifically here but the inference behavior is similar - i.e. connecting input to output (specifically, the use case where I ran into this is a profunctor).

Would it be more feasible to implement a new ad-hoc build function that is able to connect two adjacent expressions in this way? It would allow some interesting use cases to move from fluent-style API to builder-style API. The specific example I have is a "transform" type, i.e. an operator chain modeled as a value independently of the values the chain may be applied to.