   let freshName = count(from: 0)
 
   func freshVar(_ type: Type) -> AbstractValue? {
-    switch removeOwnership(type) {
+    switch simplifyType(type) {
     case .namedType("Int"): return .int(.var(IntVar(freshName())))
     case .namedType("Bool"): return .bool(.var(freshBoolVar()))
     case .namedType("TensorShape"): return freshShapeValue()

        updates = [.function(name)]

      case let .partialApply(_, _, fn, _, args, fnType):
        guard case let .functionType(allArgTypes, _) = unwrapFunctionType(fnType) else {
          fatalError("Expected a function type in .partialApply, got: \(fnType)")
        }
        assert(allArgTypes.count >= args.count)

  }
}

fileprivate func simplifyType(_ type: Type) -> Type {
  switch type {
  case let .attributedType(_, t): return simplifyType(t)
  case let .genericType(_, _, t): return simplifyType(t)
  case let .withOwnership(_, subtype): return simplifyType(subtype)
  default: return type
  }
}
