[AutoDiff] Add control flow AD leak checking tests.

[dan-zheng]

No description provided.

[dan-zheng]

Oops - this test actually has a leak count of 0.
But I recall encountering a test like this that fails with a leak count of 1.

cc @rxwei @pschuh: do you recall any reproducers for the known closure capture related memory leak?

[rxwei]

Looks like it will leak. Try wrapping things in a loop?

for _ in 0..<10 {
  let x: Tracked<Float> = 1.0
  _ = model.gradient { m in m.applied(to: x) }
}

[dan-zheng]

Added more tests in d4691b8.

Loop actually doesn't leak, but I found a test case that does leak.

  // TODO: Fix memory leak.
  testWithLeakChecking(expectedLeakCount: 1) {
    var model = ExampleLeakModel()
    let x: Tracked<Float> = 1.0
    _ = model.gradient { m in
      var model = m
      // Next line causes leak.
      model.bias = x
      return model.applied(to: x)
    }
  }

Full dump here, in case you have time to look into it along with control flow related memory leaks.

[rxwei]

Interesting. This doesn't seem related to closure checking then. What if you wrote:

    _ = model.gradient { m in
      let x: Tracked<Float> = 1.0
      var model = m
      // Next line causes leak.
      model.bias = x
      return model.applied(to: x)
    }

[dan-zheng]

Interesting. This doesn't seem related to closure checking then. What if you wrote: ...

The partial snippet you shared also produces a leak (it's pretty similar to the leak that's checked in). Here's the full snippet:

  // TODO: Fix memory leak.
  testWithLeakChecking(expectedLeakCount: 1) {
    var model = ExampleLeakModel()
    _ = model.gradient { m in
      let x: Tracked<Float> = 1.0
      var model = m
      // Next line causes leak.
      model.bias = x
      return model.applied(to: x)
    }
  }

[rxwei]

Yeah, so it should be unrelated to closure captures. It's a mutation differentiation bug. Could you file a bug for this specific issue with a reproducer?

[dan-zheng]

Filed TF-550.

[pschuh]

I think dan's example here of adding tests with expectedLeakCount is a better way to keep track of the leaks and their reproducers. Checking the leaks in would keep track of when they are fixed better. The bug could also contain the reproducer.

[dan-zheng]

@swift-ci Please test tensorflow