 // model TF's "global seed".
 @usableFromInline @inline(__always)
 func _tensorSeeds(_ seed: Tensor<Int64>) -> (Tensor<Int64>, Tensor<Int64>) {
-  return (Tensor(_defaultGraphSeed), seed)
+    return (Tensor(_defaultGraphSeed), seed)
 }
 
 //===----------------------------------------------------------------------===//

/// element tensors.
@_fixed_layout
public struct Dataset<Element : TensorGroup> {
  public let _handle: VariantHandle

  @inlinable
  public init(_handle: VariantHandle) {
    self._handle = _handle
  }
}

public extension Dataset {
  @inlinable
  init(randomSeed: Int64) {
    let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
    self.init(_handle: Raw.experimentalRandomDataset(
      seed: seed1,
      seed2: seed2,
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList))
  }
}

public extension Dataset {
  /// Creates a dataset from a batch of elements as a tensor.
  @inlinable
  init(elements: Element) {
    self.init(_handle: Raw.tensorSliceDataset(
      components: [elements],
      outputShapes: Element._unknownShapeList))
  }
}

extension Dataset : Sequence {
  public typealias Iterator = DatasetIterator<Element>

  /// Returns an iterator over the elements of this dataset.
  @inlinable
  public func makeIterator() -> DatasetIterator<Element> {
    let resource = Raw.anonymousIterator(
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList)
    Raw.makeIterator(dataset: _handle, iterator: resource)
    return DatasetIterator(_handle: resource)
  }
}

public extension Dataset {
  // Note that this Dataset API implementation uses an experimental tracing
  // feature, which is not robust and does not have great diagnostics yet.
  @inlinable
  func map<ResultElement : TensorGroup>(
    _ transform: (Element) -> ResultElement
  ) -> Dataset<ResultElement> {
    return Dataset<ResultElement>(_handle: Raw.mapDataset(
      inputDataset: _handle,
      otherArguments: Tensor<Int32>(0),
      f: transform,
      outputTypes: ResultElement._typeList,
      outputShapes: ResultElement._unknownShapeList,
      useInterOpParallelism: true,
      preserveCardinality: false))
  }

  @inlinable
  func map<ResultElement : TensorGroup>(
    parallelCallCount: Int,
    _ transform: (Element) -> ResultElement
  ) -> Dataset<ResultElement> {
    return Dataset<ResultElement>(_handle: Raw.parallelMapDataset(
      inputDataset: _handle,
      otherArguments: Tensor<Int32>(0),
      numParallelCalls: Tensor<Int32>(Int32(parallelCallCount)),
      f: transform,
      outputTypes: ResultElement._typeList,
      outputShapes: ResultElement._unknownShapeList,
      useInterOpParallelism: true,
      sloppy: false,
      preserveCardinality: false))
  }

  @inlinable
  func filter(
    _ isIncluded: (Element) -> Tensor<Bool>
  ) -> Dataset {
    return Dataset(_handle: Raw.filterDataset(
      inputDataset: _handle,
      otherArguments: Tensor<Int32>(0),
      predicate: isIncluded,
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList))
  }
}

public extension Dataset {
  @inlinable
  func shuffled(
    sampleCount: Int, randomSeed: Int64
  ) -> Dataset {
    let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
    return Dataset(_handle: Raw.shuffleDataset(
      inputDataset: _handle,
      bufferSize: Tensor(Int64(sampleCount)),
      seed: seed1,
      seed2: seed2,
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList))
  }

  @inlinable
  func batched(_ batchSize: Int) -> Dataset {
    return Dataset(_handle: Raw.batchDataset(
      inputDataset: _handle,
      batchSize: Tensor(Int64(batchSize)),
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList))
  }
}

/// The type that allows iteration over a dataset's elements.
@_fixed_layout
public struct DatasetIterator<Element : TensorGroup> {
  @usableFromInline let _handle: ResourceHandle

  @usableFromInline
  internal init(_handle: ResourceHandle) {
    self._handle = _handle
  }
}

extension DatasetIterator : IteratorProtocol {
  /// Advances to the next element and returns it, or `nil` if no next element
  /// exists.
  @inlinable
  public mutating func next() -> Element? {
    let optional = Raw.iteratorGetNextAsOptional(
      iterator: _handle,
      outputTypes: Element._typeList,
      outputShapes: Element._unknownShapeList)
    guard Raw.optionalHasValue(optional: optional).scalarized() else {
      return nil
    }
    return Raw.optionalGetValue(
      optional: optional,
      outputShapes: Element._unknownShapeList)
  }
}

/// A 2-tuple-like struct that conforms to TensorGroup that represents a tuple 
/// of 2 types conforming to TensorGroup.
@_fixed_layout
public struct Zip2TensorGroup<T : TensorGroup, U : TensorGroup> : TensorGroup {
  public var first: T
  public var second: U

  public init(_ first: T, _ second: U) {
    self.first = first
    self.second = second
  }
}

@inlinable
public func zip<T : TensorGroup, U : TensorGroup>(
  _ dataset1: Dataset<T>, _ dataset2: Dataset<U>
) -> Dataset<Zip2TensorGroup<T, U>> {
  let handle = Raw.zipDataset(
    inputDatasets: [dataset1._handle, dataset2._handle],
    outputTypes: Zip2TensorGroup<T, U>._typeList,
    outputShapes: Zip2TensorGroup<T, U>._unknownShapeList)
  return Dataset(_handle: handle)
}

// Copyright 2018 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// The default graph seed.
///
