[SR-12682] Function builder generic parameter type inference fails

[dan-zheng]

Previous ID	SR-12682
Radar	rdar://problem/62481590
Original Reporter	@dan-zheng
Type	Bug
Status	Closed
Resolution	Won't Do

Additional Detail from JIRA

| | | |------------------|-----------------| |Votes | 0 | |Component/s | Compiler | |Labels | Bug, FunctionBuilder, TypeChecker | |Assignee | None | |Priority | Medium | md5: 56625d74937947d4f1c1a26936e26eef

Issue Description:

Function builder generic parameter type inference fails, even though type inference succeeds for manually-written buildBlock calls.

Cross-posted here from https://forums.swift.org/t/function-builder-cannot-infer-generic-parameters-even-though-direct-call-to-buildblock-can/35886.

Context

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 currently 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>(...)
}

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

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

But it doesn't work:

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>

Direct calls to LayerBuilder.buildBlock do type-check inferring some generic parameters, which makes me feel that 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())

Full reproducer

Uncomment FIXME comment for type inference 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())

[beccadax]

@swift-ci create

[dan-zheng]

Forum replies (by @rjmccall) suggest that the "lack of type inference" is expected behavior, given how function builder syntax sugar is expanded.

That makes sense as a technical explanation. I'm still curious about answers to Anandabits (JIRA User)'s use-case-motivated questions:

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?

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>
}

[rjmccall]

Serious discussions on function builders' future design evolution should probably stay on the evolution forums, just to avoid having a million separate conversations that people feel obliged to follow. Please tag me when you post.

[dan-zheng]

Thanks John.

I've just replied on the forums: https://forums.swift.org/t/function-builder-cannot-infer-generic-parameters-even-though-direct-call-to-buildblock-can/35886/15.

Feel free to close this issue if centralizing discussion on the forums makes sense. I'm hesitant to close myself since there's an attached rdar.

[hborla]

Closing as behaves correctly. Follow up discussions about function builder semantics are happening on the forums.