Skip to content
This repository was archived by the owner on Apr 23, 2025. It is now read-only.

Commit 644f16b

Browse files
BradLarsonBradLarson
authored
Stock toolchain cleanup (#724)
* Remove explicit Numerics imports, enable AnyLayer to build. * Removing the last of the conditional code, temporarily disabling some mismatching AnyLayer exponentials.
1 parent fda0214 commit 644f16b

File tree

8 files changed

+198
-247
lines changed

8 files changed

+198
-247
lines changed

‎CMakeLists.txt

-21
Original file line numberDiff line numberDiff line change
@@ -148,27 +148,6 @@ set_target_properties(ArgumentParser PROPERTIES
148148
add_dependencies(ArgumentParser swift-argument-parser-install)
149149
add_dependencies(turbojpeg libjpeg-turbo-build)
150150

151-
option(ENABLE_SWIFT_NUMERICS
152-
"Enable integrating swift-numerics" YES)
153-
154-
include(CheckSwiftSourceCompiles)
155-
check_swift_source_compiles("struct S : KeyPathIterable { }" SWIFT_COMPILER_HAS_KEYPATHITERABLE_PROTOCOL)
156-
if(SWIFT_COMPILER_HAS_KEYPATHITERABLE_PROTOCOL)
157-
set(TENSORFLOW_USE_STANDARD_TOOLCHAIN_DEFAULT OFF)
158-
else()
159-
# if(CMAKE_Swift_COMPILER_VERSION VERSION_GREATER 5.3)
160-
message(WARNING "Swift compiler does not support KeyPathIterable protocol - assuming stock toolchain")
161-
set(TENSORFLOW_USE_STANDARD_TOOLCHAIN_DEFAULT ON)
162-
# endif()
163-
endif()
164-
165-
include(CMakeDependentOption)
166-
cmake_dependent_option(TENSORFLOW_USE_STANDARD_TOOLCHAIN
167-
"Experimental support to use a standard toolchain"
168-
${TENSORFLOW_USE_STANDARD_TOOLCHAIN_DEFAULT}
169-
"ENABLE_SWIFT_NUMERICS"
170-
NO)
171-
172151
add_subdirectory(Autoencoder)
173152
add_subdirectory(TensorBoard)
174153
add_subdirectory(Support)

‎Models/Text/CMakeLists.txt

-3
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,11 @@ add_library(TextModels
1616
WordSeg/SemiRing.swift)
1717
set_target_properties(TextModels PROPERTIES
1818
INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
19-
target_compile_definitions(TextModels PRIVATE
20-
$<$<BOOL:${TENSORFLOW_USE_STANDARD_TOOLCHAIN}>:TENSORFLOW_USE_STANDARD_TOOLCHAIN>)
2119
target_compile_options(TextModels PRIVATE
2220
$<$<BOOL:${BUILD_TESTING}>:-enable-testing>)
2321
target_link_libraries(TextModels PUBLIC
2422
Checkpoints
2523
Datasets
26-
$<$<AND:$<BOOL:${ENABLE_SWIFT_NUMERICS}>,$<BOOL:${TENSORFLOW_USE_STANDARD_TOOLCHAIN}>>:Numerics>
2724
SwiftProtobuf)
2825

2926
install(TARGETS TextModels

‎Models/Text/WordSeg/Model.swift

-4
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,6 @@
2323
import ModelSupport
2424
import TensorFlow
2525

26-
#if TENSORFLOW_USE_STANDARD_TOOLCHAIN
27-
import Numerics
28-
#endif
29-
3026
/// Types that can be optimized by an optimizer.
3127
///
3228
/// TODO: Consider promoting this into a public protocol in swift-apis?

‎Package.swift

+2-5
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ let package = Package(
2727
.package(url: "https://github.com/apple/swift-protobuf.git", from: "1.10.0"),
2828
.package(url: "https://github.com/apple/swift-argument-parser", .branch("main")),
2929
.package(url: "https://github.com/google/swift-benchmark", from: "0.1.0"),
30-
.package(url: "https://github.com/apple/swift-numerics", .branch("main")),
3130
],
3231
targets: [
3332
.target(
@@ -36,14 +35,12 @@ let package = Package(
3635
.target(name: "Datasets", dependencies: ["ModelSupport"], path: "Datasets"),
3736
.target(name: "STBImage", path: "Support/STBImage"),
3837
.target(
39-
name: "ModelSupport",
40-
dependencies: ["STBImage", .product(name: "Numerics", package: "swift-numerics"),],
41-
path: "Support", exclude: ["STBImage"]),
38+
name: "ModelSupport", dependencies: ["STBImage"], path: "Support", exclude: ["STBImage"]),
4239
.target(name: "TensorBoard", dependencies: ["SwiftProtobuf", "ModelSupport", "TrainingLoop"], path: "TensorBoard"),
4340
.target(name: "ImageClassificationModels", path: "Models/ImageClassification"),
4441
.target(name: "VideoClassificationModels", path: "Models/Spatiotemporal"),
4542
.target(name: "TextModels",
46-
dependencies: ["Checkpoints", "Datasets", "SwiftProtobuf", .product(name: "Numerics", package: "swift-numerics")],
43+
dependencies: ["Checkpoints", "Datasets", "SwiftProtobuf"],
4744
path: "Models/Text"),
4845
.target(name: "RecommendationModels", path: "Models/Recommendation"),
4946
.target(name: "TrainingLoop", dependencies: ["ModelSupport"], path: "TrainingLoop"),

‎Support/AnyLayer.swift

+27-30
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
// limitations under the License.
1414

1515
// TODO: Re-enable this for the stock toolchain when it can be realigned with VectorProtocol.
16-
#if !TENSORFLOW_USE_STANDARD_TOOLCHAIN
1716
import TensorFlow
1817
import _Differentiation
1918

@@ -41,27 +40,27 @@ fileprivate func mustOverride(function: StaticString = #function, file: StaticSt
4140
/// - Input: the input type of the underlying layar
4241
/// - Output: the output type of the underlying layer
4342
/// - Scalar: the scalar type of the underlying tangent vector
44-
internal class AnyLayerBox<Input: Differentiable, Output: Differentiable, Scalar: FloatingPoint & ElementaryFunctions> {
43+
internal class AnyLayerBox<Input: Differentiable, Output: Differentiable> {
4544
/// The underlying layer, type-erased to `Any`.
4645
var typeErasedBase: Any {
4746
mustOverride()
4847
}
4948

5049
/// Returns the underlying layer unboxed to the given type, if possible.
5150
func unboxed<U: Layer>(to type: U.Type) -> U?
52-
where U.TangentVector.VectorSpaceScalar == Scalar {
51+
where U.TangentVector.VectorSpaceScalar == Float {
5352
mustOverride()
5453
}
5554

5655
// `Differentiable` requirements.
5756
/// Moves `self` along the given direction. In Riemannian geometry, this is equivalent to exponential map, which moves `self` on the geodesic surface along the given tangent vector.
58-
func _move(along direction: AnyLayerTangentVector<Scalar>) {
57+
func _move(along direction: AnyLayerTangentVector) {
5958
mustOverride()
6059
}
6160

6261
// `EuclideanDifferentiable` requirements.
6362
/// The differentiable vector component of `self`.
64-
var _differentiableVectorView: AnyLayerTangentVector<Scalar> {
63+
var _differentiableVectorView: AnyLayerTangentVector {
6564
mustOverride()
6665
}
6766

@@ -72,7 +71,7 @@ internal class AnyLayerBox<Input: Differentiable, Output: Differentiable, Scalar
7271
}
7372

7473
func _vjpCallAsFunction(_ input: Input) ->
75-
(value: Output, pullback: (Output.TangentVector) -> (AnyLayerTangentVector<Scalar>, Input.TangentVector)) {
74+
(value: Output, pullback: (Output.TangentVector) -> (AnyLayerTangentVector, Input.TangentVector)) {
7675
mustOverride()
7776
}
7877

@@ -84,14 +83,14 @@ internal class AnyLayerBox<Input: Differentiable, Output: Differentiable, Scalar
8483
}
8584

8685
/// Creates a new box storing a copy of the underlying layer, used to preserve value semantics.
87-
func duplicate() -> AnyLayerBox<Input, Output, Scalar> {
86+
func duplicate() -> AnyLayerBox<Input, Output> {
8887
mustOverride()
8988
}
9089
}
9190

9291
/// A concrete implementation of the type-erased layer wrapper that forwards to an underlying layer.
93-
internal class ConcreteLayerBox<Underlying: Layer>: AnyLayerBox<Underlying.Input, Underlying.Output, Underlying.TangentVector.VectorSpaceScalar>
94-
where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunctions {
92+
internal class ConcreteLayerBox<Underlying: Layer>: AnyLayerBox<Underlying.Input, Underlying.Output>
93+
where Underlying.TangentVector.VectorSpaceScalar == Float {
9594
/// The underlying layer.
9695
var underlying: Underlying
9796

@@ -107,12 +106,12 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
107106

108107
/// Returns the underlying layer unboxed to the given type, if possible.
109108
override func unboxed<U: Layer>(to type: U.Type) -> U?
110-
where U.TangentVector.VectorSpaceScalar == Underlying.TangentVector.VectorSpaceScalar {
109+
where U.TangentVector.VectorSpaceScalar == Float {
111110
return (self as? ConcreteLayerBox<U>)?.underlying
112111
}
113112

114113
// `Differentiable` requirements.
115-
override func _move(along direction: AnyLayerTangentVector<Underlying.TangentVector.VectorSpaceScalar>) {
114+
override func _move(along direction: AnyLayerTangentVector) {
116115
if let scalarDirection = direction.box.getOpaqueScalar() {
117116
underlying.move(along: Underlying.TangentVector.zero.adding(scalarDirection))
118117
} else {
@@ -125,7 +124,7 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
125124
}
126125

127126
// `EuclideanDifferentiable` requirements.
128-
public override var _differentiableVectorView: AnyLayerTangentVector<Underlying.TangentVector.VectorSpaceScalar> {
127+
public override var _differentiableVectorView: AnyLayerTangentVector {
129128
return AnyLayerTangentVector(underlying.differentiableVectorView)
130129
}
131130

@@ -143,7 +142,7 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
143142
override func _vjpCallAsFunction(_ input: Underlying.Input) -> (
144143
value: Underlying.Output,
145144
pullback: (Underlying.Output.TangentVector) ->
146-
(AnyLayerTangentVector<Underlying.TangentVector.VectorSpaceScalar>, Underlying.Input.TangentVector)
145+
(AnyLayerTangentVector, Underlying.Input.TangentVector)
147146
) {
148147
let basePullback = valueWithPullback(
149148
at: ModelAndInput(model: underlying, input: input),
@@ -155,7 +154,7 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
155154
pullback: { (outTangent) in
156155
let pairTangent = basePullback.pullback(outTangent)
157156
return (
158-
AnyLayerTangentVector<Underlying.TangentVector.VectorSpaceScalar>(pairTangent.model),
157+
AnyLayerTangentVector(pairTangent.model),
159158
pairTangent.input
160159
)
161160
}
@@ -164,12 +163,12 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
164163

165164
// `CopyableToDevice` requirements.
166165
override func _copyToDevice(to device: Device) ->
167-
AnyLayerBox<Underlying.Input, Underlying.Output, Underlying.TangentVector.VectorSpaceScalar> {
166+
AnyLayerBox<Underlying.Input, Underlying.Output> {
168167
return ConcreteLayerBox(Underlying(copying: underlying, to: device))
169168
}
170169

171170
override func duplicate() ->
172-
AnyLayerBox<Underlying.Input, Underlying.Output, Underlying.TangentVector.VectorSpaceScalar> {
171+
AnyLayerBox<Underlying.Input, Underlying.Output> {
173172
return ConcreteLayerBox(underlying)
174173
}
175174
}
@@ -189,11 +188,10 @@ where Underlying.TangentVector.VectorSpaceScalar: FloatingPoint & ElementaryFunc
189188
/// Type Parameters:
190189
/// - Input: the input type of the underlying layar
191190
/// - Output: the output type of the underlying layer
192-
/// - Scalar: the scalar type of the underlying tangent vector
193-
public struct AnyLayer<Input: Differentiable, Output: Differentiable, Scalar: FloatingPoint & ElementaryFunctions>: CopyableToDevice {
194-
internal var box: AnyLayerBox<Input, Output, Scalar>
191+
public struct AnyLayer<Input: Differentiable, Output: Differentiable>: CopyableToDevice {
192+
internal var box: AnyLayerBox<Input, Output>
195193

196-
internal init(box: AnyLayerBox<Input, Output, Scalar>) {
194+
internal init(box: AnyLayerBox<Input, Output>) {
197195
self.box = box
198196
}
199197

@@ -205,7 +203,7 @@ public struct AnyLayer<Input: Differentiable, Output: Differentiable, Scalar: Fl
205203
/// Creates a type-erased derivative from the given layer.
206204
@differentiable
207205
public init<Underlying: Layer>(_ layer: Underlying)
208-
where Underlying.Input == Input, Underlying.Output == Output, Underlying.TangentVector.VectorSpaceScalar == Scalar {
206+
where Underlying.Input == Input, Underlying.Output == Output, Underlying.TangentVector.VectorSpaceScalar == Float {
209207
self.box = ConcreteLayerBox<Underlying>(layer)
210208
}
211209

@@ -217,25 +215,25 @@ public struct AnyLayer<Input: Differentiable, Output: Differentiable, Scalar: Fl
217215
@derivative(of: init)
218216
internal static func _vjpInit<T: Layer>(
219217
_ base: T
220-
) -> (value: AnyLayer, pullback: (AnyLayerTangentVector<Scalar>) -> T.TangentVector)
221-
where T.Input == Input, T.Output == Output, T.TangentVector.VectorSpaceScalar == Scalar
218+
) -> (value: AnyLayer, pullback: (AnyLayerTangentVector) -> T.TangentVector)
219+
where T.Input == Input, T.Output == Output, T.TangentVector.VectorSpaceScalar == Float
222220
{
223-
return (AnyLayer<Input, Output, Scalar>(base), { v in v.unboxed(as: T.TangentVector.self)! })
221+
return (AnyLayer<Input, Output>(base), { v in v.unboxed(as: T.TangentVector.self)! })
224222
}
225223

226224
@inlinable
227225
@derivative(of: init)
228226
internal static func _jvpInit<T: Layer>(
229227
_ base: T
230228
) -> (
231-
value: AnyLayer, differential: (T.TangentVector) -> AnyLayerTangentVector<Scalar>
232-
) where T.Input == Input, T.Output == Output, T.TangentVector.VectorSpaceScalar == Scalar {
233-
return (AnyLayer<Input, Output, Scalar>(base), { dbase in AnyLayerTangentVector<Scalar>(dbase) })
229+
value: AnyLayer, differential: (T.TangentVector) -> AnyLayerTangentVector
230+
) where T.Input == Input, T.Output == Output, T.TangentVector.VectorSpaceScalar == Float {
231+
return (AnyLayer<Input, Output>(base), { dbase in AnyLayerTangentVector(dbase) })
234232
}
235233
}
236234

237235
extension AnyLayer: Differentiable {
238-
public typealias TangentVector = AnyLayerTangentVector<Scalar>
236+
public typealias TangentVector = AnyLayerTangentVector
239237

240238
public mutating func move(along direction: TangentVector) {
241239
if !isKnownUniquelyReferenced(&box) { // preserve value semantics
@@ -260,7 +258,7 @@ extension AnyLayer: Layer {
260258

261259
@derivative(of: _callAsFunction)
262260
func _vjpCallAsFunction(_ input: Input) ->
263-
(value: Output, pullback: (Output.TangentVector) -> (AnyLayerTangentVector<Scalar>, Input.TangentVector)) {
261+
(value: Output, pullback: (Output.TangentVector) -> (AnyLayerTangentVector, Input.TangentVector)) {
264262
return box._vjpCallAsFunction(input)
265263
}
266264

@@ -269,4 +267,3 @@ extension AnyLayer: Layer {
269267
return _callAsFunction(input)
270268
}
271269
}
272-
#endif

0 commit comments

Comments
 (0)