Upload 161 files

allocation_description.proto

@@ -0,0 +1,27 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "AllocationDescriptionProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+message AllocationDescription {
+  // Total number of bytes requested
+  int64 requested_bytes = 1;
+
+  // Total number of bytes allocated if known
+  int64 allocated_bytes = 2;
+
+  // Name of the allocator used
+  string allocator_name = 3;
+
+  // Identifier of the allocated buffer if known
+  int64 allocation_id = 4;
+
+  // Set if this tensor only has one remaining reference
+  bool has_single_reference = 5;
+
+  // Address of the allocation.
+  uint64 ptr = 6;
+};

allocator.cc

@@ -0,0 +1,130 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/allocator.h"
+
+#include "tensorflow/core/framework/allocator_registry.h"
+#include "tensorflow/core/framework/log_memory.h"
+#include "tensorflow/core/framework/tracking_allocator.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
+#include "tensorflow/core/platform/mem.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+void AllocatorStats::Clear() {
+  this->num_allocs = 0;
+  this->bytes_in_use = 0;
+  this->max_bytes_in_use = 0;
+  this->max_alloc_size = 0;
+  this->bytes_limit = 0;
+}
+
+string AllocatorStats::DebugString() const {
+  return strings::Printf(
+      "Limit:        %20lld\n"
+      "InUse:        %20lld\n"
+      "MaxInUse:     %20lld\n"
+      "NumAllocs:    %20lld\n"
+      "MaxAllocSize: %20lld\n",
+      this->bytes_limit, this->bytes_in_use, this->max_bytes_in_use,
+      this->num_allocs, this->max_alloc_size);
+}
+
+constexpr size_t Allocator::kAllocatorAlignment;
+
+Allocator::~Allocator() {}
+
+void RunResourceCtor(ResourceHandle* p, size_t n) {
+  for (size_t i = 0; i < n; ++p, ++i) new (p) ResourceHandle();
+}
+
+void RunResourceDtor(ResourceHandle* p, size_t n) {
+  for (size_t i = 0; i < n; ++p, ++i) p->~ResourceHandle();
+}
+
+// If true, cpu allocator collects more stats.
+static bool cpu_allocator_collect_stats = false;
+// If true, cpu allocator collects full stats.
+static bool cpu_allocator_collect_full_stats = false;
+
+void EnableCPUAllocatorStats(bool enable) {
+  cpu_allocator_collect_stats = enable;
+}
+void EnableCPUAllocatorFullStats(bool enable) {
+  cpu_allocator_collect_full_stats = enable;
+}
+
+class CPUAllocator : public Allocator {
+ public:
+  CPUAllocator() {}
+
+  ~CPUAllocator() override {}
+
+  string Name() override { return "cpu"; }
+
+  void* AllocateRaw(size_t alignment, size_t num_bytes) override {
+    void* p = port::AlignedMalloc(num_bytes, alignment);
+    if (cpu_allocator_collect_stats) {
+      const std::size_t alloc_size = port::MallocExtension_GetAllocatedSize(p);
+      mutex_lock l(mu_);
+      ++stats_.num_allocs;
+      stats_.bytes_in_use += alloc_size;
+      stats_.max_bytes_in_use =
+          std::max<int64>(stats_.max_bytes_in_use, stats_.bytes_in_use);
+      stats_.max_alloc_size =
+          std::max<int64>(stats_.max_alloc_size, alloc_size);
+    }
+    return p;
+  }
+
+  void DeallocateRaw(void* ptr) override {
+    if (cpu_allocator_collect_stats) {
+      const std::size_t alloc_size =
+          port::MallocExtension_GetAllocatedSize(ptr);
+      mutex_lock l(mu_);
+      stats_.bytes_in_use -= alloc_size;
+    }
+    port::AlignedFree(ptr);
+  }
+
+  void GetStats(AllocatorStats* stats) override {
+    mutex_lock l(mu_);
+    *stats = stats_;
+  }
+
+  size_t AllocatedSizeSlow(void* ptr) override {
+    return port::MallocExtension_GetAllocatedSize(ptr);
+  }
+
+ private:
+  mutex mu_;
+  AllocatorStats stats_ GUARDED_BY(mu_);
+
+  TF_DISALLOW_COPY_AND_ASSIGN(CPUAllocator);
+};
+
+Allocator* cpu_allocator() {
+  static Allocator* cpu_alloc = AllocatorRegistry::Global()->GetAllocator();
+  if (cpu_allocator_collect_full_stats && !cpu_alloc->TracksAllocationSizes()) {
+    cpu_alloc = new TrackingAllocator(cpu_alloc, true);
+  }
+  return cpu_alloc;
+}
+
+REGISTER_MEM_ALLOCATOR("DefaultCPUAllocator", 100, CPUAllocator);
+
+}  // namespace tensorflow

allocator.h

@@ -0,0 +1,394 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_ALLOCATOR_H_
+#define TENSORFLOW_FRAMEWORK_ALLOCATOR_H_
+
+#include <stdlib.h>
+
+#include <limits>
+
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/type_traits.h"
+#include "tensorflow/core/framework/variant.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Attributes for a single allocation call. Different calls to the same
+// allocator could potentially have different allocation attributes.
+struct AllocationAttributes {
+  // If the first attempt to allocate the memory fails, the allocation
+  // should return immediately without retrying.
+  // An example use case is optional scratch spaces where a failure
+  // has only performance impact.
+  bool no_retry_on_failure = false;
+  // If a Tensor is allocated without the following set to true, then
+  // it is logged as an unknown allocation. During execution Tensors
+  // should be allocated through the OpKernelContext which records
+  // which Op is performing the allocation, and sets this flag to
+  // true.
+  bool allocation_will_be_logged = false;
+};
+
+// Runtime statistics collected by an allocator.
+struct AllocatorStats {
+  int64 num_allocs;        // Number of allocations.
+  int64 bytes_in_use;      // Number of bytes in use.
+  int64 max_bytes_in_use;  // The maximum bytes in use.
+  int64 max_alloc_size;    // The max single allocation seen.
+
+  // The upper limit what the allocator can allocate, if such a limit
+  // is known. Certain allocator may return 0 to indicate the limit is
+  // unknown.
+  int64 bytes_limit;
+
+  AllocatorStats() { Clear(); }
+
+  void Clear();
+  string DebugString() const;
+};
+
+// Allocator is an abstract interface for allocating and deallocating
+// device memory.
+class Allocator {
+ public:
+#ifdef EIGEN_VECTORIZE_AVX512
+  // Align to 64 byte boundary.
+  static constexpr size_t kAllocatorAlignment = 64;
+#else
+  // Align to 32 byte boundary.
+  static constexpr size_t kAllocatorAlignment = 32;
+#endif
+
+  virtual ~Allocator();
+
+  // Return a string identifying this allocator
+  virtual string Name() = 0;
+
+  // Return an uninitialized block of memory that is "num_bytes" bytes
+  // in size.  The returned pointer is guaranteed to be aligned to a
+  // multiple of "alignment" bytes.
+  // REQUIRES: "alignment" is a power of 2.
+  virtual void* AllocateRaw(size_t alignment, size_t num_bytes) = 0;
+
+  // Return an uninitialized block of memory that is "num_bytes" bytes
+  // in size with specified allocation attributes.  The returned pointer is
+  // guaranteed to be aligned to a multiple of "alignment" bytes.
+  // REQUIRES: "alignment" is a power of 2.
+  virtual void* AllocateRaw(size_t alignment, size_t num_bytes,
+                            const AllocationAttributes& allocation_attr) {
+    // The default behavior is to use the implementation without any allocation
+    // attributes.
+    return AllocateRaw(alignment, num_bytes);
+  }
+
+  // Deallocate a block of memory pointer to by "ptr"
+  // REQUIRES: "ptr" was previously returned by a call to AllocateRaw
+  virtual void DeallocateRaw(void* ptr) = 0;
+
+  // Convenience functions to do typed allocation.  C++ constructors
+  // and destructors are invoked for complex types if necessary,
+  // depending on the concrete Allocator implementation. May return
+  // NULL if the tensor has too many elements to represent in a single
+  // allocation.
+  template <typename T>
+  T* Allocate(size_t num_elements) {
+    return Allocate<T>(num_elements, AllocationAttributes());
+  }
+
+  template <typename T>
+  T* Allocate(size_t num_elements,
+              const AllocationAttributes& allocation_attr) {
+    // TODO(jeff): Do we need to allow clients to pass in alignment
+    // requirements?
+
+    if (num_elements > (std::numeric_limits<size_t>::max() / sizeof(T))) {
+      return NULL;
+    }
+
+    void* p = AllocateRaw(kAllocatorAlignment, sizeof(T) * num_elements,
+                          allocation_attr);
+    T* typed_p = reinterpret_cast<T*>(p);
+    if (typed_p) RunCtor<T>(typed_p, num_elements);
+    return typed_p;
+  }
+
+  template <typename T>
+  void Deallocate(T* ptr, size_t num_elements) {
+    if (ptr) {
+      RunDtor<T>(ptr, num_elements);
+      DeallocateRaw(ptr);
+    }
+  }
+
+  // Returns true if this allocator tracks the sizes of allocations.
+  // RequestedSize and AllocatedSize must be overridden if
+  // TracksAllocationSizes is overridden to return true.
+  virtual bool TracksAllocationSizes() { return false; }
+
+  // Returns true if this allocator requires tensors with 0 elements
+  // to allocate buffers. This is false for most allocators, but may
+  // be used by special-case allocators that want to track tensor
+  // usage.
+  virtual bool ShouldAllocateEmptyTensors() { return false; }
+
+  // Returns the user-requested size of the data allocated at
+  // 'ptr'.  Note that the actual buffer allocated might be larger
+  // than requested, but this function returns the size requested by
+  // the user.
+  //
+  // REQUIRES: TracksAllocationSizes() is true.
+  //
+  // REQUIRES: 'ptr!=nullptr' and points to a buffer previously
+  // allocated by this allocator.
+  virtual size_t RequestedSize(void* ptr) {
+    CHECK(false) << "allocator doesn't track sizes";
+    return size_t(0);
+  }
+
+  // Returns the allocated size of the buffer at 'ptr' if known,
+  // otherwise returns RequestedSize(ptr). AllocatedSize(ptr) is
+  // guaranteed to be >= RequestedSize(ptr).
+  //
+  // REQUIRES: TracksAllocationSizes() is true.
+  //
+  // REQUIRES: 'ptr!=nullptr' and points to a buffer previously
+  // allocated by this allocator.
+  virtual size_t AllocatedSize(void* ptr) { return RequestedSize(ptr); }
+
+  // Returns either 0 or an identifier assigned to the buffer at 'ptr'
+  // when the buffer was returned by AllocateRaw. If non-zero, the
+  // identifier differs from every other ID assigned by this
+  // allocator.
+  //
+  // REQUIRES: TracksAllocationSizes() is true.
+  //
+  // REQUIRES: 'ptr!=nullptr' and points to a buffer previously
+  // allocated by this allocator.
+  virtual int64 AllocationId(void* ptr) { return 0; }
+
+  // Returns the allocated size of the buffer at 'ptr' if known,
+  // otherwise returns 0. This method can be called when
+  // TracksAllocationSizes() is false, but can be extremely slow.
+  //
+  // REQUIRES: 'ptr!=nullptr' and points to a buffer previously
+  // allocated by this allocator.
+  virtual size_t AllocatedSizeSlow(void* ptr) {
+    if (TracksAllocationSizes()) {
+      return AllocatedSize(ptr);
+    }
+    return 0;
+  }
+
+  // Fills in 'stats' with statistics collected by this allocator.
+  virtual void GetStats(AllocatorStats* stats) { stats->Clear(); }
+
+ private:
+  // No constructors or destructors are run for simple types
+  template <typename T>
+  void RunCtor(T* p, size_t n) {
+    static_assert(is_simple_type<T>::value, "T is not a simple type.");
+  }
+
+  template <typename T>
+  void RunDtor(T* p, size_t n) {}
+
+  // custom constructors and destructors that can be overridden for
+  // non-standard allocators
+
+  // Runs string's default constructor for  p[0], p[1], ..., p[n-1].
+  virtual void RunStringCtor(string* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) new (p) string();
+  }
+
+  // Runs string's default destructor for  p[0], p[1], ..., p[n-1].
+  virtual void RunStringDtor(string* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) p->~string();
+  }
+
+  virtual void RunResourceCtor(ResourceHandle* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) new (p) ResourceHandle();
+  }
+
+  // Runs string's default destructor for  p[0], p[1], ..., p[n-1].
+  virtual void RunResourceDtor(ResourceHandle* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) p->~ResourceHandle();
+  }
+
+  virtual void RunVariantCtor(Variant* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) new (p) Variant();
+  }
+
+  virtual void RunVariantDtor(Variant* p, size_t n) {
+    for (size_t i = 0; i < n; ++p, ++i) p->~Variant();
+  }
+
+  // TODO(jeff): Maybe provide some interface to give info about
+  // current allocation state (total number of bytes available for
+  // allocation, number of bytes free on device, etc.)
+};
+
+// Allocator-specific constructors and destructors are used for
+// strings
+template <>
+inline void Allocator::RunCtor(string* p, size_t n) {
+  RunStringCtor(p, n);
+}
+
+template <>
+inline void Allocator::RunDtor(string* p, size_t n) {
+  RunStringDtor(p, n);
+}
+
+template <>
+inline void Allocator::RunCtor(ResourceHandle* p, size_t n) {
+  RunResourceCtor(p, n);
+}
+
+template <>
+inline void Allocator::RunDtor(ResourceHandle* p, size_t n) {
+  RunResourceDtor(p, n);
+}
+
+template <>
+inline void Allocator::RunCtor(Variant* p, size_t n) {
+  RunVariantCtor(p, n);
+}
+
+template <>
+inline void Allocator::RunDtor(Variant* p, size_t n) {
+  RunVariantDtor(p, n);
+}
+
+// An implementation of Allocator that delegates all calls to another Allocator.
+//
+// Useful to clients who want to override part of the functionality of another
+// allocator.
+class AllocatorWrapper : public Allocator {
+ public:
+  explicit AllocatorWrapper(Allocator* wrapped) : wrapped_(wrapped) {}
+
+  ~AllocatorWrapper() override {}
+
+  // Returns the wrapped allocator to which all calls are delegated.
+  Allocator* wrapped() const { return wrapped_; }
+
+  string Name() override { return wrapped_->Name(); }
+
+  void* AllocateRaw(size_t alignment, size_t num_bytes) override {
+    return wrapped_->AllocateRaw(alignment, num_bytes);
+  }
+
+  void* AllocateRaw(size_t alignment, size_t num_bytes,
+                    const AllocationAttributes& allocation_attr) override {
+    return wrapped_->AllocateRaw(alignment, num_bytes, allocation_attr);
+  }
+
+  void DeallocateRaw(void* ptr) override { wrapped_->DeallocateRaw(ptr); }
+
+  bool TracksAllocationSizes() override {
+    return wrapped_->TracksAllocationSizes();
+  }
+
+  bool ShouldAllocateEmptyTensors() override {
+    return wrapped_->TracksAllocationSizes();
+  }
+
+  size_t RequestedSize(void* ptr) override {
+    return wrapped_->RequestedSize(ptr);
+  }
+
+  size_t AllocatedSize(void* ptr) override {
+    return wrapped_->AllocatedSize(ptr);
+  }
+
+  int64 AllocationId(void* ptr) override { return wrapped_->AllocationId(ptr); }
+
+  size_t AllocatedSizeSlow(void* ptr) override {
+    return wrapped_->AllocatedSizeSlow(ptr);
+  }
+
+ private:
+  Allocator* const wrapped_;
+};
+
+// A tensorflow Op may need access to different kinds of memory that
+// are not simply a function of the device to which the Op has been
+// assigned.  For example, an Op executing on a GPU may still need
+// to allocate CPU RAM for some purpose.  Internal to the tensorflow
+// runtime we may choose to allocate CPU ram from special regions
+// that have been prepared for higher performance in some use
+// contexts, e.g. doing DMA with particular devices.  For these
+// reasons, the Device interface does not expose just one memory
+// Allocator, but instead provides an accessor that takes a
+// specification of the desired memory attributes in order to select
+// an Allocator.
+//
+// Example use:
+//  // Allocator for ordinary device memory:
+//  Allocator* a = allocator(AllocatorAttributes());
+// ...
+//  // Allocator for CPU RAM, regardless of where Op is executing:
+//  AllocatorAttributes attr;
+//  attr.set_on_host(true);
+//  Allocator* a = allocator(attr);
+struct AllocatorAttributes {
+  void set_on_host(bool v) { value |= (static_cast<int>(v)); }
+  bool on_host() const { return value & 0x1; }
+  void set_nic_compatible(bool v) { value |= (static_cast<int>(v) << 1); }
+  bool nic_compatible() const { return value & (0x1 << 1); }
+  void set_gpu_compatible(bool v) { value |= (static_cast<int>(v) << 2); }
+  bool gpu_compatible() const { return value & (0x1 << 2); }
+  void Merge(AllocatorAttributes other) { value |= other.value; }
+  // Returns true if the fields set in *this is a subset of or equal to
+  // those set in other.
+  bool IsEqualOrLessRestrictiveThan(const AllocatorAttributes& other) const {
+    return (value | other.value) == other.value;
+  }
+
+  // NOTE: The upper 8 bits of the value are reserved for
+  // device-specific uses.  Implementors of a device can interpret these
+  // upper 8 bits in device-specific ways, and ops implemented for those
+  // devices are responsible for setting those 8 bits appropriately.
+  uint32 value = 0;
+};
+
+// Returns a trivial implementation of Allocator which uses the system
+// default malloc. The returned allocator is a process singleton.
+Allocator* cpu_allocator();
+
+// If 'enable' is true, the process-wide cpu allocator collects
+// AllocatorStats. By default, it's disabled.
+void EnableCPUAllocatorStats(bool enable);
+
+// If 'enable' is true, the process-wide cpu allocator collects full
+// statistics. By default, it's disabled.
+void EnableCPUAllocatorFullStats(bool enable);
+
+// Abstract interface of an object that does the underlying suballoc/free of
+// memory for a higher-level allocator.
+class SubAllocator {
+ public:
+  virtual ~SubAllocator() {}
+  virtual void* Alloc(size_t alignment, size_t num_bytes) = 0;
+  virtual void Free(void* ptr, size_t num_bytes) = 0;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_ALLOCATOR_H_

allocator_registry.cc

@@ -0,0 +1,80 @@
+/* Copyright 2017 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.
+==============================================================================*/
+
+#include <string>
+
+#include "tensorflow/core/framework/allocator_registry.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+// static
+AllocatorRegistry* AllocatorRegistry::Global() {
+  static AllocatorRegistry* global_allocator_registry = new AllocatorRegistry;
+  return global_allocator_registry;
+}
+
+Allocator* AllocatorRegistry::GetRegisteredAllocator(const string& name,
+                                                     int priority) {
+  for (auto entry : allocators_) {
+    if (!name.compare(entry.name) && priority == entry.priority) {
+      return entry.allocator;
+    }
+  }
+  return nullptr;
+}
+
+void AllocatorRegistry::Register(const string& name, int priority,
+                                 Allocator* allocator) {
+  CHECK(!name.empty()) << "Need a valid name for Allocator";
+  CHECK_GE(priority, 0) << "Priority needs to be non-negative";
+
+  Allocator* existing = GetRegisteredAllocator(name, priority);
+  if (existing != nullptr) {
+    // A duplicate is if the registration name and priority match
+    // but the Allocator::Name()'s don't match.
+    CHECK_EQ(existing->Name(), allocator->Name())
+        << "Allocator with name: [" << name << "], type [" << existing->Name()
+        << "], priority: [" << priority
+        << "] already registered.  Choose a different name to register "
+        << "an allocator of type " << allocator->Name();
+
+    // The allocator names match, so we can just return.
+    // It should be safe to delete the allocator since the caller
+    // gives up ownership of it.
+    delete allocator;
+    return;
+  }
+
+  AllocatorRegistryEntry tmp_entry;
+  tmp_entry.name = name;
+  tmp_entry.priority = priority;
+  tmp_entry.allocator = allocator;
+
+  allocators_.push_back(tmp_entry);
+  int high_pri = -1;
+  for (auto entry : allocators_) {
+    if (high_pri < entry.priority) {
+      m_curr_allocator_ = entry.allocator;
+      high_pri = entry.priority;
+    }
+  }
+}
+
+Allocator* AllocatorRegistry::GetAllocator() {
+  return CHECK_NOTNULL(m_curr_allocator_);
+}
+
+}  // namespace tensorflow

allocator_registry.h

@@ -0,0 +1,80 @@
+/* Copyright 2017 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.
+==============================================================================*/
+
+// Classes to maintain a static registry of memory allocators
+#ifndef TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_
+#define TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/allocator.h"
+
+namespace tensorflow {
+
+// A global AllocatorRegistry is used to hold allocators for CPU backends
+class AllocatorRegistry {
+ public:
+  // Add an allocator to the registry.  Caller releases ownership of
+  // 'allocator'.
+  void Register(const string& name, int priority, Allocator* allocator);
+
+  // Return allocator with highest priority
+  // If multiple allocators have the same high priority, return one of them
+  Allocator* GetAllocator();
+
+  // Returns the global registry of allocators.
+  static AllocatorRegistry* Global();
+
+ private:
+  typedef struct {
+    string name;
+    int priority;
+    Allocator* allocator;  // not owned
+  } AllocatorRegistryEntry;
+
+  // Returns the Allocator registered for 'name' and 'priority',
+  // or 'nullptr' if not found.
+  Allocator* GetRegisteredAllocator(const string& name, int priority);
+
+  std::vector<AllocatorRegistryEntry> allocators_;
+  Allocator* m_curr_allocator_;  // not owned
+};
+
+namespace allocator_registration {
+
+class AllocatorRegistration {
+ public:
+  AllocatorRegistration(const string& name, int priority,
+                        Allocator* allocator) {
+    AllocatorRegistry::Global()->Register(name, priority, allocator);
+  }
+};
+
+}  // namespace allocator_registration
+
+#define REGISTER_MEM_ALLOCATOR(name, priority, allocator) \
+  REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(__COUNTER__, name, priority, allocator)
+
+#define REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(ctr, name, priority, allocator) \
+  REGISTER_MEM_ALLOCATOR_UNIQ(ctr, name, priority, allocator)
+
+#define REGISTER_MEM_ALLOCATOR_UNIQ(ctr, name, priority, allocator) \
+  static allocator_registration::AllocatorRegistration              \
+      register_allocator_##ctr(name, priority, new allocator)
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_

allocator_test.cc

@@ -0,0 +1,186 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/allocator.h"
+
+#include <algorithm>
+#include <vector>
+
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+
+static void CheckStats(Allocator* a, int64 num_allocs, int64 bytes_in_use,
+                       int64 max_bytes_in_use, int64 max_alloc_size) {
+  AllocatorStats stats;
+  a->GetStats(&stats);
+  LOG(INFO) << "Alloc stats: \n" << stats.DebugString();
+#if defined(PLATFORM_GOOGLE) && defined(NDEBUG)
+  // NOTE: allocator stats expectation depends on the system malloc,
+  // and can vary as that changes.
+  static const int64 kSlop = 5 * 1024;
+  EXPECT_GT(stats.bytes_in_use, bytes_in_use - kSlop);
+  EXPECT_LT(stats.bytes_in_use, bytes_in_use + kSlop);
+  EXPECT_GT(stats.max_bytes_in_use, max_bytes_in_use - kSlop);
+  EXPECT_LT(stats.max_bytes_in_use, max_bytes_in_use + kSlop);
+  EXPECT_EQ(stats.num_allocs, num_allocs);
+  EXPECT_EQ(stats.max_alloc_size, max_alloc_size);
+#endif
+}
+
+TEST(AllocatorAttributesTest, AllCombos) {
+  for (bool on_host : {false, true}) {
+    for (bool nic_compatible : {false, true}) {
+      for (bool gpu_compatible : {false, true}) {
+        AllocatorAttributes aa;
+        aa.set_on_host(on_host);
+        aa.set_nic_compatible(nic_compatible);
+        aa.set_gpu_compatible(gpu_compatible);
+        EXPECT_EQ(on_host, aa.on_host());
+        EXPECT_EQ(nic_compatible, aa.nic_compatible());
+        EXPECT_EQ(gpu_compatible, aa.gpu_compatible());
+      }
+    }
+  }
+}
+
+TEST(AllocatorAttributesTest, IsEqualOrLessRestrictiveThan) {
+  AllocatorAttributes a, b;
+  EXPECT_TRUE(a.IsEqualOrLessRestrictiveThan(b));
+  EXPECT_TRUE(a.IsEqualOrLessRestrictiveThan(a));
+  EXPECT_TRUE(b.IsEqualOrLessRestrictiveThan(b));
+
+  b.set_gpu_compatible(true);
+  // The set of flags in b is not a subset of those in a.
+  EXPECT_TRUE(a.IsEqualOrLessRestrictiveThan(b));
+  EXPECT_FALSE(b.IsEqualOrLessRestrictiveThan(a));
+  EXPECT_TRUE(a.IsEqualOrLessRestrictiveThan(a));
+  EXPECT_TRUE(b.IsEqualOrLessRestrictiveThan(b));
+
+  a.set_nic_compatible(true);
+  // Neither a nor b is a subset of the other.
+  EXPECT_FALSE(a.IsEqualOrLessRestrictiveThan(b));
+  EXPECT_FALSE(b.IsEqualOrLessRestrictiveThan(a));
+
+  a.set_gpu_compatible(true);
+  // The set of flags in b is a proper subset of those in a.
+  EXPECT_TRUE(b.IsEqualOrLessRestrictiveThan(a));
+  EXPECT_FALSE(a.IsEqualOrLessRestrictiveThan(b));
+}
+
+TEST(CPUAllocatorTest, Simple) {
+  EnableCPUAllocatorStats(true);
+  Allocator* a = cpu_allocator();
+  std::vector<void*> ptrs;
+  for (int s = 1; s < 1024; s++) {
+    void* raw = a->AllocateRaw(1, s);
+    ptrs.push_back(raw);
+  }
+  std::sort(ptrs.begin(), ptrs.end());
+  CheckStats(a, 1023, 552640, 552640, 1024);
+  for (size_t i = 0; i < ptrs.size(); i++) {
+    if (i > 0) {
+      CHECK_NE(ptrs[i], ptrs[i - 1]);  // No dups
+    }
+    a->DeallocateRaw(ptrs[i]);
+  }
+  CheckStats(a, 1023, 0, 552640, 1024);
+  float* t1 = a->Allocate<float>(1024);
+  double* t2 = a->Allocate<double>(1048576);
+  CheckStats(a, 1025, 1048576 * sizeof(double) + 1024 * sizeof(float),
+             1048576 * sizeof(double) + 1024 * sizeof(float),
+             1048576 * sizeof(double));
+
+  a->Deallocate(t1, 1024);
+  a->Deallocate(t2, 1048576);
+
+  CheckStats(a, 1025, 0, 1048576 * sizeof(double) + 1024 * sizeof(float),
+             1048576 * sizeof(double));
+  EnableCPUAllocatorStats(false);
+}
+
+// Define a struct that we will use to observe behavior in the unit tests
+struct TestStruct {
+  int x;  // not used just want to make sure sizeof(TestStruct) > 1
+};
+
+TEST(CPUAllocatorTest, CheckStructSize) { CHECK_GT(sizeof(TestStruct), 1); }
+
+TEST(CPUAllocatorTest, AllocateOverflowMaxSizeT) {
+  Allocator* a = cpu_allocator();
+
+  // The maximum size_t value will definitely overflow.
+  size_t count_to_allocate = std::numeric_limits<size_t>::max();
+  TestStruct* const test_pointer = a->Allocate<TestStruct>(count_to_allocate);
+
+  CHECK_EQ(test_pointer, reinterpret_cast<TestStruct*>(NULL));
+}
+
+TEST(CPUAllocatorTest, AllocateOverflowSmallest) {
+  Allocator* a = cpu_allocator();
+
+  // count_to_allocate is the smallest count that will cause overflow.
+  const size_t count_to_allocate =
+      (std::numeric_limits<size_t>::max() / sizeof(TestStruct)) + 1;
+  TestStruct* const test_pointer = a->Allocate<TestStruct>(count_to_allocate);
+
+  CHECK_EQ(test_pointer, reinterpret_cast<TestStruct*>(NULL));
+}
+
+TEST(CPUAllocatorTest, Sizes) {
+  Allocator* a = cpu_allocator();
+
+  EXPECT_EQ(false, a->TracksAllocationSizes());
+}
+
+namespace {
+
+AllocatorAttributes DeviceAllocatorAttribute() {
+  AllocatorAttributes attr;
+  attr.value |= (0x1 << 24);
+  return attr;
+}
+
+bool HasDeviceAllocatorAttribute(const AllocatorAttributes& attr) {
+  return attr.value & (0x1 << 24);
+}
+
+}  // namespace
+
+TEST(CustomAllocatorAttributes, TestSetterAndGetter) {
+  AllocatorAttributes attr = DeviceAllocatorAttribute();
+  EXPECT_TRUE(HasDeviceAllocatorAttribute(attr));
+  EXPECT_FALSE(HasDeviceAllocatorAttribute(AllocatorAttributes()));
+}
+
+static void BM_Allocation(int iters, int arg) {
+  Allocator* a = cpu_allocator();
+  // Exercise a few different allocation sizes
+  std::vector<int> sizes = {256, 4096, 16384, 524288, 512, 1048576};
+  int size_index = 0;
+
+  if (arg) EnableCPUAllocatorStats(true);
+  while (--iters > 0) {
+    int bytes = sizes[size_index++ % sizes.size()];
+    void* p = a->AllocateRaw(1, bytes);
+    a->DeallocateRaw(p);
+  }
+  if (arg) EnableCPUAllocatorStats(false);
+}
+BENCHMARK(BM_Allocation)->Arg(0)->Arg(1);
+
+}  // namespace tensorflow

api_def.proto

@@ -0,0 +1,120 @@
+// Defines the text format for including per-op API definition and
+// overrides for client language op code generators.
+
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "ApiDefProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+import "tensorflow/core/framework/attr_value.proto";
+
+// Used to specify and override the default API & behavior in the
+// generated code for client languages, from what you would get from
+// the OpDef alone. There will be a set of ApiDefs that are common
+// to all client languages, and another set per client language.
+// The per-client-language ApiDefs will inherit values from the
+// common ApiDefs which it can either replace or modify.
+//
+// We separate the API definition from the OpDef so we can evolve the
+// API while remaining backwards compatible when interpretting old
+// graphs.  Overrides go in an "api_def.pbtxt" file with a text-format
+// ApiDefs message.
+//
+// WARNING: Be *very* careful changing the API for any existing op --
+// you can change the semantics of existing code.  These changes may
+// need to wait until a major release of TensorFlow to avoid breaking
+// our compatibility promises.
+message ApiDef {
+  // Name of the op (in the OpDef) to specify the API for.
+  string graph_op_name = 1;
+
+  enum Visibility {
+    // Normally this is "VISIBLE" unless you are inheriting a
+    // different value from another ApiDef.
+    DEFAULT_VISIBILITY = 0;
+    // Publicly visible in the API.
+    VISIBLE = 1;
+    // Do not include this op in the generated API. If visibility is
+    // set to 'SKIP', other fields are ignored for this op.
+    SKIP = 2;
+    // Hide this op by putting it into an internal namespace (or whatever
+    // is appropriate in the target language).
+    HIDDEN = 3;
+  }
+  Visibility visibility = 2;
+
+  // If you specify any endpoint, this will replace all of the
+  // inherited endpoints.  The first endpoint should be the
+  // "canonical" endpoint, and should not be deprecated (unless all
+  // endpoints are deprecated).
+  message Endpoint {
+    // Name should be either like "CamelCaseName" or
+    // "Package.CamelCaseName". Client-language-specific ApiDefs may
+    // use a snake_case convention instead of CamelCase.
+    string name = 1;
+
+    // First GraphDef version at which the op is disallowed.
+    int32 deprecation_version = 2;
+  }
+  repeated Endpoint endpoint = 3;
+
+  message Arg {
+    string name = 1;
+
+    // Change the name used to access this arg in the API from what
+    // is used in the GraphDef.  Note that these names in `backticks`
+    // will also be replaced in the summary & description fields.
+    string rename_to = 2;
+
+    // Note: this will replace any inherited arg doc. There is no
+    // current way of modifying arg descriptions (other than replacing
+    // them entirely) as can be done with op descriptions.
+    string description = 3;
+  }
+  repeated Arg in_arg = 4;
+  repeated Arg out_arg = 5;
+  // List of original in_arg names to specify new argument order.
+  // Length of arg_order should be either empty to keep current order
+  // or match size of in_arg.
+  repeated string arg_order = 11;
+
+  // Description of the graph-construction-time configuration of this
+  // Op.  That is to say, this describes the attr fields that will
+  // be specified in the NodeDef.
+  message Attr {
+    string name = 1;
+
+    // Change the name used to access this attr in the API from what
+    // is used in the GraphDef.  Note that these names in `backticks`
+    // will also be replaced in the summary & description fields.
+    string rename_to = 2;
+
+    // Specify a new default value to use for this attr.  This default
+    // will be used when creating new graphs, as opposed to the
+    // default in the OpDef, which will be used when interpreting old
+    // GraphDefs.
+    AttrValue default_value = 3;
+
+    // Note: this will replace any inherited attr doc, there is no current
+    // way of modifying attr descriptions as can be done with op descriptions.
+    string description = 4;
+  }
+  repeated Attr attr = 6;
+
+  // One-line human-readable description of what the Op does.
+  string summary = 7;
+
+  // Additional, longer human-readable description of what the Op does.
+  string description = 8;
+
+  // Modify an existing/inherited description by adding text to the beginning
+  // or end.
+  string description_prefix = 9;
+  string description_suffix = 10;
+}
+
+message ApiDefs {
+  repeated ApiDef op = 1;
+}

attr_value.proto

@@ -0,0 +1,62 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "AttrValueProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/tensor.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
+
+// Protocol buffer representing the value for an attr used to configure an Op.
+// Comment indicates the corresponding attr type.  Only the field matching the
+// attr type may be filled.
+message AttrValue {
+  // LINT.IfChange
+  message ListValue {
+    repeated bytes s = 2;                        // "list(string)"
+    repeated int64 i = 3 [packed = true];        // "list(int)"
+    repeated float f = 4 [packed = true];        // "list(float)"
+    repeated bool b = 5 [packed = true];         // "list(bool)"
+    repeated DataType type = 6 [packed = true];  // "list(type)"
+    repeated TensorShapeProto shape = 7;         // "list(shape)"
+    repeated TensorProto tensor = 8;             // "list(tensor)"
+    repeated NameAttrList func = 9;              // "list(attr)"
+  }
+  // LINT.ThenChange(https://www.tensorflow.org/code/tensorflow/c/c_api.cc)
+
+  oneof value {
+    bytes s = 2;                 // "string"
+    int64 i = 3;                 // "int"
+    float f = 4;                 // "float"
+    bool b = 5;                  // "bool"
+    DataType type = 6;           // "type"
+    TensorShapeProto shape = 7;  // "shape"
+    TensorProto tensor = 8;      // "tensor"
+    ListValue list = 1;          // any "list(...)"
+
+    // "func" represents a function. func.name is a function's name or
+    // a primitive op's name. func.attr.first is the name of an attr
+    // defined for that function. func.attr.second is the value for
+    // that attr in the instantiation.
+    NameAttrList func = 10;
+
+    // This is a placeholder only used in nodes defined inside a
+    // function.  It indicates the attr value will be supplied when
+    // the function is instantiated.  For example, let us suppose a
+    // node "N" in function "FN". "N" has an attr "A" with value
+    // placeholder = "foo". When FN is instantiated with attr "foo"
+    // set to "bar", the instantiated node N's attr A will have been
+    // given the value "bar".
+    string placeholder = 9;
+  }
+}
+
+// A list of attr names and their values. The whole list is attached
+// with a string name.  E.g., MatMul[T=float].
+message NameAttrList {
+  string name = 1;
+  map<string, AttrValue> attr = 2;
+}

attr_value_util.cc

@@ -0,0 +1,551 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/attr_value_util.h"
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/attr_value.pb_text.h"
+#include "tensorflow/core/framework/tensor.pb_text.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb_text.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+namespace {
+
+string SummarizeString(const string& str) {
+  return strings::StrCat("\"", str_util::CEscape(str), "\"");
+}
+
+string SummarizeTensor(const TensorProto& tensor_proto) {
+  Tensor t;
+  if (!t.FromProto(tensor_proto)) {
+    return strings::StrCat(
+        "<Invalid TensorProto: ", ProtoShortDebugString(tensor_proto), ">");
+  }
+  return t.DebugString();
+}
+
+string SummarizeFunc(const NameAttrList& func) {
+  std::vector<string> entries;
+  for (auto p : func.attr()) {
+    entries.push_back(
+        strings::StrCat(p.first, "=", SummarizeAttrValue(p.second)));
+  }
+  std::sort(entries.begin(), entries.end());
+  return strings::StrCat(func.name(), "[", str_util::Join(entries, ", "), "]");
+}
+
+}  // namespace
+
+string SummarizeAttrValue(const AttrValue& attr_value) {
+  switch (attr_value.value_case()) {
+    case AttrValue::kS:
+      return SummarizeString(attr_value.s());
+    case AttrValue::kI:
+      return strings::StrCat(attr_value.i());
+    case AttrValue::kF:
+      return strings::StrCat(attr_value.f());
+    case AttrValue::kB:
+      return attr_value.b() ? "true" : "false";
+    case AttrValue::kType:
+      return EnumName_DataType(attr_value.type());
+    case AttrValue::kShape:
+      return PartialTensorShape::DebugString(attr_value.shape());
+    case AttrValue::kTensor:
+      return SummarizeTensor(attr_value.tensor());
+    case AttrValue::kList: {
+      string ret = "[";
+      if (attr_value.list().s_size() > 0) {
+        for (int i = 0; i < attr_value.list().s_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret, SummarizeString(attr_value.list().s(i)));
+        }
+      } else if (attr_value.list().i_size() > 0) {
+        for (int i = 0; i < attr_value.list().i_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret, attr_value.list().i(i));
+        }
+      } else if (attr_value.list().f_size() > 0) {
+        for (int i = 0; i < attr_value.list().f_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret, attr_value.list().f(i));
+        }
+      } else if (attr_value.list().b_size() > 0) {
+        for (int i = 0; i < attr_value.list().b_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret, attr_value.list().b(i) ? "true" : "false");
+        }
+      } else if (attr_value.list().type_size() > 0) {
+        for (int i = 0; i < attr_value.list().type_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret,
+                             EnumName_DataType(attr_value.list().type(i)));
+        }
+      } else if (attr_value.list().shape_size() > 0) {
+        for (int i = 0; i < attr_value.list().shape_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(
+              &ret, TensorShape::DebugString(attr_value.list().shape(i)));
+        }
+      } else if (attr_value.list().tensor_size() > 0) {
+        for (int i = 0; i < attr_value.list().tensor_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret,
+                             SummarizeTensor(attr_value.list().tensor(i)));
+        }
+      } else if (attr_value.list().func_size() > 0) {
+        for (int i = 0; i < attr_value.list().func_size(); ++i) {
+          if (i > 0) strings::StrAppend(&ret, ", ");
+          strings::StrAppend(&ret, SummarizeFunc(attr_value.list().func(i)));
+        }
+      }
+
+      strings::StrAppend(&ret, "]");
+      return ret;
+    }
+    case AttrValue::kFunc: {
+      return SummarizeFunc(attr_value.func());
+    }
+    case AttrValue::kPlaceholder:
+      return strings::StrCat("$", attr_value.placeholder());
+    case AttrValue::VALUE_NOT_SET:
+      return "<Unknown AttrValue type>";
+  }
+  return "<Unknown AttrValue type>";  // Prevent missing return warning
+}
+
+Status AttrValueHasType(const AttrValue& attr_value, StringPiece type) {
+  int num_set = 0;
+
+#define VALIDATE_FIELD(name, type_string, oneof_case)                         \
+  do {                                                                        \
+    if (attr_value.has_list()) {                                              \
+      if (attr_value.list().name##_size() > 0) {                              \
+        if (type != "list(" type_string ")") {                                \
+          return errors::InvalidArgument(                                     \
+              "AttrValue had value with type 'list(" type_string ")' when '", \
+              type, "' expected");                                            \
+        }                                                                     \
+        ++num_set;                                                            \
+      }                                                                       \
+    } else if (attr_value.value_case() == AttrValue::oneof_case) {            \
+      if (type != type_string) {                                              \
+        return errors::InvalidArgument(                                       \
+            "AttrValue had value with type '" type_string "' when '", type,   \
+            "' expected");                                                    \
+      }                                                                       \
+      ++num_set;                                                              \
+    }                                                                         \
+  } while (false)
+
+  VALIDATE_FIELD(s, "string", kS);
+  VALIDATE_FIELD(i, "int", kI);
+  VALIDATE_FIELD(f, "float", kF);
+  VALIDATE_FIELD(b, "bool", kB);
+  VALIDATE_FIELD(type, "type", kType);
+  VALIDATE_FIELD(shape, "shape", kShape);
+  VALIDATE_FIELD(tensor, "tensor", kTensor);
+  VALIDATE_FIELD(func, "func", kFunc);
+
+#undef VALIDATE_FIELD
+
+  if (attr_value.value_case() == AttrValue::kPlaceholder) {
+    return errors::InvalidArgument(
+        "AttrValue had value with unexpected type 'placeholder'");
+  }
+
+  // If the attr type is 'list', we expect attr_value.has_list() to be
+  // true.  However, proto3's attr_value.has_list() can be false when
+  // set to an empty list for GraphDef versions <= 4. So we simply
+  // check if has_list is false and some other field in attr_value is
+  // set to flag the error.  This test can be made more strict once
+  // support for GraphDef versions <= 4 is dropped.
+  if (StringPiece(type).starts_with("list(") && !attr_value.has_list()) {
+    if (num_set) {
+      return errors::InvalidArgument(
+          "AttrValue missing value with expected type '", type, "'");
+    } else {
+      // Indicate that we have a list, but an empty one.
+      ++num_set;
+    }
+  }
+
+  // Okay to have an empty list, but not to be missing a non-list value.
+  if (num_set == 0 && !StringPiece(type).starts_with("list(")) {
+    return errors::InvalidArgument(
+        "AttrValue missing value with expected type '", type, "'");
+  }
+
+  // Ref types and DT_INVALID are illegal, and DataTypes must
+  // be a valid enum type.
+  if (type == "type") {
+    if (!DataType_IsValid(attr_value.type())) {
+      return errors::InvalidArgument("AttrValue has invalid DataType enum: ",
+                                     attr_value.type());
+    }
+    if (IsRefType(attr_value.type())) {
+      return errors::InvalidArgument(
+          "AttrValue must not have reference type value of ",
+          DataTypeString(attr_value.type()));
+    }
+    if (attr_value.type() == DT_INVALID) {
+      return errors::InvalidArgument("AttrValue has invalid DataType");
+    }
+  } else if (type == "list(type)") {
+    for (auto as_int : attr_value.list().type()) {
+      const DataType dtype = static_cast<DataType>(as_int);
+      if (!DataType_IsValid(dtype)) {
+        return errors::InvalidArgument("AttrValue has invalid DataType enum: ",
+                                       as_int);
+      }
+      if (IsRefType(dtype)) {
+        return errors::InvalidArgument(
+            "AttrValue must not have reference type value of ",
+            DataTypeString(dtype));
+      }
+      if (dtype == DT_INVALID) {
+        return errors::InvalidArgument("AttrValue contains invalid DataType");
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+bool ParseAttrValue(StringPiece type, StringPiece text, AttrValue* out) {
+  // Parse type.
+  string field_name;
+  bool is_list = type.Consume("list(");
+  if (type.Consume("string")) {
+    field_name = "s";
+  } else if (type.Consume("int")) {
+    field_name = "i";
+  } else if (type.Consume("float")) {
+    field_name = "f";
+  } else if (type.Consume("bool")) {
+    field_name = "b";
+  } else if (type.Consume("type")) {
+    field_name = "type";
+  } else if (type.Consume("shape")) {
+    field_name = "shape";
+  } else if (type.Consume("tensor")) {
+    field_name = "tensor";
+  } else if (type.Consume("func")) {
+    field_name = "func";
+  } else if (type.Consume("placeholder")) {
+    field_name = "placeholder";
+  } else {
+    return false;
+  }
+  if (is_list && !type.Consume(")")) {
+    return false;
+  }
+
+  // Construct a valid text proto message to parse.
+  string to_parse;
+  if (is_list) {
+    // TextFormat parser considers "i: 7" to be the same as "i: [7]",
+    // but we only want to allow list values with [].
+    StringPiece cleaned = text;
+    str_util::RemoveLeadingWhitespace(&cleaned);
+    str_util::RemoveTrailingWhitespace(&cleaned);
+    if (cleaned.size() < 2 || cleaned[0] != '[' ||
+        cleaned[cleaned.size() - 1] != ']') {
+      return false;
+    }
+    cleaned.remove_prefix(1);
+    str_util::RemoveLeadingWhitespace(&cleaned);
+    if (cleaned.size() == 1) {
+      // User wrote "[]", so return empty list without invoking the TextFormat
+      // parse which returns an error for "i: []".
+      out->Clear();
+      out->mutable_list();
+      return true;
+    }
+    to_parse = strings::StrCat("list { ", field_name, ": ", text, " }");
+  } else {
+    to_parse = strings::StrCat(field_name, ": ", text);
+  }
+
+  return ProtoParseFromString(to_parse, out);
+}
+
+void SetAttrValue(const AttrValue& value, AttrValue* out) { *out = value; }
+
+#define DEFINE_SET_ATTR_VALUE_ONE(ARG_TYPE, FIELD) \
+  void SetAttrValue(ARG_TYPE value, AttrValue* out) { out->set_##FIELD(value); }
+
+#define DEFINE_SET_ATTR_VALUE_LIST(ARG_TYPE, FIELD)                       \
+  void SetAttrValue(ARG_TYPE value, AttrValue* out) {                     \
+    out->mutable_list()->Clear(); /* create list() even if value empty */ \
+    for (const auto& v : value) {                                         \
+      out->mutable_list()->add_##FIELD(v);                                \
+    }                                                                     \
+  }
+
+#define DEFINE_SET_ATTR_VALUE_BOTH(ARG_TYPE, FIELD) \
+  DEFINE_SET_ATTR_VALUE_ONE(ARG_TYPE, FIELD)        \
+  DEFINE_SET_ATTR_VALUE_LIST(gtl::ArraySlice<ARG_TYPE>, FIELD)
+
+DEFINE_SET_ATTR_VALUE_ONE(const string&, s)
+DEFINE_SET_ATTR_VALUE_LIST(gtl::ArraySlice<string>, s)
+DEFINE_SET_ATTR_VALUE_BOTH(const char*, s)
+DEFINE_SET_ATTR_VALUE_BOTH(int64, i)
+DEFINE_SET_ATTR_VALUE_BOTH(int32, i)
+DEFINE_SET_ATTR_VALUE_BOTH(float, f)
+DEFINE_SET_ATTR_VALUE_BOTH(double, f)
+DEFINE_SET_ATTR_VALUE_BOTH(bool, b)
+DEFINE_SET_ATTR_VALUE_LIST(const std::vector<bool>&, b)
+DEFINE_SET_ATTR_VALUE_LIST(std::initializer_list<bool>, b)
+DEFINE_SET_ATTR_VALUE_BOTH(DataType, type)
+
+void SetAttrValue(StringPiece value, AttrValue* out) {
+  out->set_s(value.data(), value.size());
+}
+
+void SetAttrValue(const gtl::ArraySlice<StringPiece> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    out->mutable_list()->add_s(v.data(), v.size());
+  }
+}
+
+void SetAttrValue(const TensorShape& value, AttrValue* out) {
+  value.AsProto(out->mutable_shape());
+}
+
+void SetAttrValue(const TensorShapeProto& value, AttrValue* out) {
+  *out->mutable_shape() = value;
+}
+
+void SetAttrValue(const PartialTensorShape& value, AttrValue* out) {
+  value.AsProto(out->mutable_shape());
+}
+
+void SetAttrValue(const gtl::ArraySlice<TensorShape> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    v.AsProto(out->mutable_list()->add_shape());
+  }
+}
+
+void SetAttrValue(gtl::ArraySlice<TensorShapeProto> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    *out->mutable_list()->add_shape() = v;
+  }
+}
+
+void SetAttrValue(const gtl::ArraySlice<PartialTensorShape> value,
+                  AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    v.AsProto(out->mutable_list()->add_shape());
+  }
+}
+
+void SetAttrValue(const Tensor& value, AttrValue* out) {
+  if (value.NumElements() > 1) {
+    value.AsProtoTensorContent(out->mutable_tensor());
+  } else {
+    value.AsProtoField(out->mutable_tensor());
+  }
+}
+
+void SetAttrValue(const gtl::ArraySlice<Tensor> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    if (v.NumElements() > 1) {
+      v.AsProtoTensorContent(out->mutable_list()->add_tensor());
+    } else {
+      v.AsProtoField(out->mutable_list()->add_tensor());
+    }
+  }
+}
+
+void SetAttrValue(const TensorProto& value, AttrValue* out) {
+  *out->mutable_tensor() = value;
+}
+
+void SetAttrValue(const gtl::ArraySlice<TensorProto> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    *out->mutable_list()->add_tensor() = v;
+  }
+}
+
+void SetAttrValue(const NameAttrList& value, AttrValue* out) {
+  *out->mutable_func() = value;
+}
+
+void SetAttrValue(gtl::ArraySlice<NameAttrList> value, AttrValue* out) {
+  out->mutable_list()->Clear();  // Create list() even if value empty.
+  for (const auto& v : value) {
+    *out->mutable_list()->add_func() = v;
+  }
+}
+
+bool AreAttrValuesEqual(const AttrValue& a, const AttrValue& b) {
+  // There are multiple equivalent representations of attr values containing
+  // TensorProtos. Compare them by constructing Tensors and serializing them
+  // back. Comparing Tensor objects is pretty tricky.
+  if (a.has_tensor() != b.has_tensor()) {
+    return false;
+  } else if (a.has_tensor() && b.has_tensor()) {
+    Tensor at(a.tensor().dtype());
+    bool success = at.FromProto(a.tensor());
+    DCHECK(success);
+
+    Tensor bt(b.tensor().dtype());
+    success = bt.FromProto(b.tensor());
+    DCHECK(success);
+
+    TensorProto ap;
+    at.AsProtoTensorContent(&ap);
+
+    TensorProto bp;
+    bt.AsProtoTensorContent(&bp);
+
+    string a_str, b_str;
+    SerializeToStringDeterministic(ap, &a_str);
+    SerializeToStringDeterministic(bp, &b_str);
+    return a_str == b_str;
+  }
+
+  // `func` field contains a nested AttrValue. Compare such AttrValues
+  // recursively.
+  if (a.has_func() != b.has_func()) {
+    return false;
+  } else if (a.has_func() && b.has_func()) {
+    const NameAttrList& af = a.func();
+    const NameAttrList& bf = b.func();
+    if (af.name() != bf.name()) return false;
+    std::unordered_map<string, AttrValue> am(af.attr().begin(),
+                                             af.attr().end());
+    for (const auto& bm_pair : bf.attr()) {
+      const auto& iter = am.find(bm_pair.first);
+      if (iter == am.end()) return false;
+      if (!AreAttrValuesEqual(iter->second, bm_pair.second)) return false;
+      am.erase(iter);
+    }
+    if (!am.empty()) return false;
+    return true;
+  }
+
+  // All other fields in AttrValue have deterministic representations.
+  // It is safe to compare their serialized strings.
+  string a_str, b_str;
+  SerializeToStringDeterministic(a, &a_str);
+  SerializeToStringDeterministic(b, &b_str);
+  return a_str == b_str;
+}
+
+uint64 AttrValueHash(const AttrValue& a) {
+  if (a.has_tensor()) {
+    // Deal with multiple representations by parsing TensorProto to
+    // Tensor and serializing it back. This is slow, but current use case
+    // don't need high efficiency.
+    Tensor tensor(a.tensor().dtype());
+    bool success = tensor.FromProto(a.tensor());
+    DCHECK(success);
+    TensorProto p;
+    tensor.AsProtoTensorContent(&p);
+    string s;
+    SerializeToStringDeterministic(p, &s);
+    return Hash64(s);
+  }
+  if (a.has_func()) {
+    const NameAttrList& func = a.func();
+    uint64 h = Hash64(func.name());
+    std::map<string, AttrValue> map(func.attr().begin(), func.attr().end());
+    for (const auto& pair : map) {
+      h = Hash64(pair.first.data(), pair.first.size(), h);
+      h = Hash64Combine(AttrValueHash(pair.second), h);
+    }
+    return h;
+  }
+
+  // If `a` is not a tensor or func, get a hash of serialized string.
+  string s;
+  SerializeToStringDeterministic(a, &s);
+  return Hash64(s);
+}
+
+bool HasPlaceHolder(const AttrValue& val) {
+  switch (val.value_case()) {
+    case AttrValue::kList: {
+      for (const NameAttrList& func : val.list().func()) {
+        for (const auto& p : func.attr()) {
+          if (HasPlaceHolder(p.second)) {
+            return true;
+          }
+        }
+      }
+      break;
+    }
+    case AttrValue::kFunc:
+      for (const auto& p : val.func().attr()) {
+        if (HasPlaceHolder(p.second)) {
+          return true;
+        }
+      }
+      break;
+    case AttrValue::kPlaceholder:
+      return true;
+    default:
+      break;
+  }
+  return false;
+}
+
+bool SubstitutePlaceholders(const SubstituteFunc& substitute,
+                            AttrValue* value) {
+  switch (value->value_case()) {
+    case AttrValue::kList: {
+      for (NameAttrList& func : *value->mutable_list()->mutable_func()) {
+        for (auto& p : *func.mutable_attr()) {
+          if (!SubstitutePlaceholders(substitute, &p.second)) {
+            return false;
+          }
+        }
+      }
+      break;
+    }
+    case AttrValue::kFunc:
+      for (auto& p : *(value->mutable_func()->mutable_attr())) {
+        if (!SubstitutePlaceholders(substitute, &p.second)) {
+          return false;
+        }
+      }
+      break;
+    case AttrValue::kPlaceholder:
+      return substitute(value->placeholder(), value);
+    case AttrValue::VALUE_NOT_SET:
+      return false;
+    default:
+      break;
+  }
+  return true;
+}
+
+}  // namespace tensorflow

attr_value_util.h

@@ -0,0 +1,116 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_ATTR_VALUE_UTIL_H_
+#define TENSORFLOW_FRAMEWORK_ATTR_VALUE_UTIL_H_
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/partial_tensor_shape.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+
+namespace tensorflow {
+
+// Forward declare protos so their symbols can be removed from .so exports
+class AttrValue;
+class NameAttrList;
+
+// A human-readable rendering of attr_value, that is more concise than a
+// text-format proto.
+string SummarizeAttrValue(const AttrValue& attr_value);
+
+// Generates an error if attr_value doesn't have the indicated attr type.
+Status AttrValueHasType(const AttrValue& attr_value, StringPiece type);
+
+// Converts a text proto value from "text" into the field of *out
+// indicated by "type" (e.g. from the type field of an AttrDef).
+// Examples:
+// * If type:"int" and text:"-14", then *out is set to "i: -14"
+// * If type:"list(string)" and text:"['foo', 'bar']",
+//   then *out is set to "list { s: ['foo', 'bar'] }"
+// Returns true on success.
+bool ParseAttrValue(StringPiece type, StringPiece text, AttrValue* out);
+
+// Sets *out based on the type of value.
+void SetAttrValue(const string& value, AttrValue* out);
+void SetAttrValue(const char* value, AttrValue* out);
+void SetAttrValue(StringPiece value, AttrValue* out);
+void SetAttrValue(int64 value, AttrValue* out);
+void SetAttrValue(int32 value, AttrValue* out);
+void SetAttrValue(float value, AttrValue* out);
+void SetAttrValue(double value, AttrValue* out);
+void SetAttrValue(bool value, AttrValue* out);
+void SetAttrValue(DataType value, AttrValue* out);
+void SetAttrValue(const TensorShape& value, AttrValue* out);
+void SetAttrValue(const TensorShapeProto& value, AttrValue* out);
+void SetAttrValue(const PartialTensorShape& value, AttrValue* out);
+void SetAttrValue(const Tensor& value, AttrValue* out);
+void SetAttrValue(const TensorProto& value, AttrValue* out);
+void SetAttrValue(const NameAttrList& value, AttrValue* out);
+
+void SetAttrValue(gtl::ArraySlice<string> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<const char*> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<StringPiece> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<int64> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<int32> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<float> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<double> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<bool> value, AttrValue* out);
+void SetAttrValue(const std::vector<bool>& value, AttrValue* out);
+void SetAttrValue(std::initializer_list<bool> value, AttrValue* out);
+void SetAttrValue(DataTypeSlice value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<TensorShape> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<TensorShapeProto> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<PartialTensorShape> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<Tensor> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<TensorProto> value, AttrValue* out);
+void SetAttrValue(gtl::ArraySlice<NameAttrList> value, AttrValue* out);
+
+void SetAttrValue(const AttrValue& value, AttrValue* out);
+
+// Returns true if a and b have the same value.
+bool AreAttrValuesEqual(const AttrValue& a, const AttrValue& b);
+
+// Returns a hash of `a` that is consistent with AreAttrValuesEqual. In other
+// words, if two AttrValues compare equal according to AreAttrValuesEqual,
+// they will have the same hash value.
+// Similarly to protobuf deterministic serialization, hash value is
+// guaranteed to be stable only for a given binary. In particular, one should
+// probably not persist the returned value.
+uint64 AttrValueHash(const AttrValue& a);
+
+// Returns true if "val" has a placeholder.
+bool HasPlaceHolder(const AttrValue& val);
+
+// SubstitutePlaceholders recursively replaces placeholders in 'value'
+// with an attr value by calling SubstituteFunc. Returns true iff all
+// placeholders in "value" are replaced with a value.
+//
+// SubstituteFunc is given a placeholder string. If the placeholder is
+// unknown, SubstituteFunc returns false. Otherwise, overwrites the
+// attr value and returns true.
+using SubstituteFunc = std::function<bool(const string&, AttrValue*)>;
+bool SubstitutePlaceholders(const SubstituteFunc& substitute, AttrValue* value);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_ATTR_VALUE_UTIL_H_

attr_value_util_test.cc

@@ -0,0 +1,195 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/attr_value_util.h"
+
+#include <vector>
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+// A few helpers to construct AttrValue protos.
+template <typename T>
+AttrValue V(T value) {
+  AttrValue ret;
+  SetAttrValue(value, &ret);
+  return ret;
+}
+
+AttrValue P(const string& p) {
+  AttrValue ret;
+  ret.set_placeholder(p);
+  return ret;
+}
+
+AttrValue F(const string& name,
+            std::vector<std::pair<string, AttrValue>> pairs) {
+  AttrValue ret;
+  ret.mutable_func()->set_name(name);
+  ret.mutable_func()->mutable_attr()->insert(pairs.begin(), pairs.end());
+  return ret;
+}
+
+AttrValue Fs(
+    std::vector<std::pair<string, std::vector<std::pair<string, AttrValue>>>>
+        funcs) {
+  AttrValue ret;
+  for (const auto& func : funcs) {
+    NameAttrList* entry = ret.mutable_list()->add_func();
+    entry->set_name(func.first);
+    entry->mutable_attr()->insert(func.second.begin(), func.second.end());
+  }
+  return ret;
+}
+
+TEST(AttrValueUtil, HasType) {
+  // OK
+  EXPECT_TRUE(AttrValueHasType(V(123), "int").ok());
+  EXPECT_TRUE(AttrValueHasType(V(1.2), "float").ok());
+  EXPECT_TRUE(AttrValueHasType(V(DT_FLOAT), "type").ok());
+  EXPECT_TRUE(AttrValueHasType(F("f", {}), "func").ok());
+  EXPECT_TRUE(AttrValueHasType(Fs({{"f", {}}, {"g", {}}}), "list(func)").ok());
+
+  // not OK.
+  EXPECT_FALSE(AttrValueHasType(V(123), "func").ok());
+  EXPECT_FALSE(AttrValueHasType(V(1.2), "int").ok());
+  EXPECT_FALSE(AttrValueHasType(V(DT_FLOAT), "shape").ok());
+  EXPECT_FALSE(AttrValueHasType(F("f", {}), "string").ok());
+  EXPECT_FALSE(AttrValueHasType(P("T"), "float").ok());
+  EXPECT_FALSE(AttrValueHasType(V(static_cast<DataType>(1000)), "type").ok());
+  std::vector<DataType> list_type({static_cast<DataType>(1000)});
+  EXPECT_FALSE(AttrValueHasType(V(list_type), "list(type)").ok());
+}
+
+SubstituteFunc ReplaceTWith(const AttrValue& val) {
+  return [val](const string& placeholder, AttrValue* target) {
+    if (placeholder == "T") {
+      *target = val;
+      return true;
+    } else {
+      return false;
+    }
+  };
+}
+
+TEST(AttrValueUtil, Basic) {
+  auto v = F("MatMul", {{"dtype", P("T")},
+                        {"transpose_a", V(false)},
+                        {"transpose_b", V(true)},
+                        {"use_cublas", V(true)}});
+  TF_EXPECT_OK(AttrValueHasType(v, "func"));
+  EXPECT_TRUE(HasPlaceHolder(v));
+
+  EXPECT_EQ(
+      SummarizeAttrValue(v),
+      "MatMul[dtype=$T, transpose_a=false, transpose_b=true, use_cublas=true]");
+
+  SubstitutePlaceholders(ReplaceTWith(V(DT_FLOAT)), &v);
+  EXPECT_TRUE(!HasPlaceHolder(v));
+  EXPECT_EQ(SummarizeAttrValue(v),
+            "MatMul[dtype=DT_FLOAT, transpose_a=false, transpose_b=true, "
+            "use_cublas=true]");
+}
+
+TEST(AttrValueUtil, Shaped) {
+  auto v =
+      F("OpRequiresShape", {{"shape_full", V(TensorShape({1, 0}))},
+                            {"shape_part", V(PartialTensorShape({-1, 1, 0}))}});
+  TF_EXPECT_OK(AttrValueHasType(v, "func"));
+  EXPECT_FALSE(HasPlaceHolder(v));
+
+  EXPECT_EQ(SummarizeAttrValue(v),
+            "OpRequiresShape[shape_full=[1,0], shape_part=[?,1,0]]");
+}
+
+TEST(AttrValueUtil, DeepAttr) {
+  auto v = Fs({{"f", {{"T", P("T")}}}, {"g", {{"T", P("T")}}}});
+  TF_EXPECT_OK(AttrValueHasType(v, "list(func)"));
+  EXPECT_TRUE(HasPlaceHolder(v));
+
+  for (int i = 0; i < 3; ++i) {
+    v = F("f", {{"T", P("T")}, {"F", v}});
+    EXPECT_TRUE(HasPlaceHolder(v));
+  }
+  EXPECT_EQ(SummarizeAttrValue(v),
+            "f[F=f[F=f[F=[f[T=$T], g[T=$T]], T=$T], T=$T], T=$T]");
+
+  SubstitutePlaceholders(ReplaceTWith(F("x", {})), &v);
+  EXPECT_TRUE(!HasPlaceHolder(v));
+  EXPECT_EQ(SummarizeAttrValue(v),
+            "f[F=f[F=f[F=[f[T=x[]], g[T=x[]]], T=x[]], T=x[]], T=x[]]");
+}
+
+AttrValue FromText(const string& text) {
+  AttrValue attr;
+  EXPECT_TRUE(protobuf::TextFormat::MergeFromString(text, &attr));
+  return attr;
+}
+
+void ExpectDifferent(const AttrValue& a1, const AttrValue& a2) {
+  EXPECT_FALSE(AreAttrValuesEqual(a1, a2));
+  EXPECT_FALSE(AreAttrValuesEqual(a2, a1));
+  EXPECT_NE(AttrValueHash(a1), AttrValueHash(a2));
+}
+
+TEST(AttrValueEquality, StringAndFuncTensors) {
+  AttrValue a = FromText(R"(
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+          dim {
+            size: 2
+          }
+        }
+        string_val: 'reader_dataset_ops_test/tmphtXHks/text_line.0.txt'
+        string_val: 'reader_dataset_ops_test/tmphtXHks/text_line.1.txt'
+      })");
+  EXPECT_TRUE(AreAttrValuesEqual(a, a));
+  EXPECT_EQ(AttrValueHash(a), AttrValueHash(a));
+
+  AttrValue b = a;
+  (*b.mutable_tensor()->mutable_string_val(0))[3] = '1';
+  ExpectDifferent(a, b);
+
+  AttrValue c1;
+  c1.mutable_func()->set_name("func_name");
+  (*c1.mutable_func()->mutable_attr())["attr1"] = a;
+  (*c1.mutable_func()->mutable_attr())["attr2"] = b;
+  EXPECT_TRUE(AreAttrValuesEqual(c1, c1));
+  EXPECT_EQ(AttrValueHash(c1), AttrValueHash(c1));
+
+  ExpectDifferent(c1, a);
+
+  AttrValue c2 = c1;
+  c2.mutable_func()->set_name("func_name2");
+  ExpectDifferent(c1, c2);
+
+  c2 = c1;
+  (*c2.mutable_func()->mutable_attr())["attr3"] = b;
+  ExpectDifferent(c1, c2);
+
+  c2 = c1;
+  (*c2.mutable_func()->mutable_attr())["attr2"] = a;
+  ExpectDifferent(c1, c2);
+
+  c2 = c1;
+  c2.mutable_func()->mutable_attr()->erase("attr2");
+  ExpectDifferent(c1, c2);
+}
+
+}  // namespace tensorflow

bfloat16.cc

@@ -0,0 +1,50 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/bfloat16.h"
+
+namespace tensorflow {
+
+void FloatToBFloat16(const float* src, bfloat16* dst, int64 size) {
+  const uint16_t* p = reinterpret_cast<const uint16_t*>(src);
+  uint16_t* q = reinterpret_cast<uint16_t*>(dst);
+#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+    for (; size != 0; p += 2, q++, size--) {
+      *q = p[0];
+    }
+#else
+    for (; size != 0; p += 2, q++, size--) {
+     *q = p[1];
+    }
+#endif
+}
+
+void BFloat16ToFloat(const bfloat16* src, float* dst, int64 size) {
+  const uint16_t* p = reinterpret_cast<const uint16_t*>(src);
+  uint16_t* q = reinterpret_cast<uint16_t*>(dst);
+#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+    for (; size != 0; p++, q += 2, size--) {
+      q[0] = *p;
+      q[1] = 0;
+    }
+#else
+    for (; size != 0; p++, q += 2, size--) {
+      q[0] = 0;
+      q[1] = *p;
+    }
+#endif
+}
+
+}  // end namespace tensorflow

bfloat16.h

@@ -0,0 +1,62 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_BFLOAT16_H_
+#define TENSORFLOW_FRAMEWORK_BFLOAT16_H_
+
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/platform/types.h"
+
+#if defined(PLATFORM_WINDOWS)
+#include "tensorflow/core/platform/windows/cpu_info.h"
+#endif
+
+// Compact 16-bit encoding of floating point numbers. This representation uses
+// 1 bit for the sign, 8 bits for the exponent and 7 bits for the mantissa.  It
+// is assumed that floats are in IEEE 754 format so the representation is just
+// bits 16-31 of a single precision float.
+//
+// NOTE: The IEEE floating point standard defines a float16 format that
+// is different than this format (it has fewer bits of exponent and more
+// bits of mantissa).  We don't use that format here because conversion
+// to/from 32-bit floats is more complex for that format, and the
+// conversion for this format is very simple.
+//
+// Because of the existing IEEE float16 type, we do not name our representation
+// "float16" but just use "uint16".
+//
+// <-----our 16bits float------->
+// s e e e e e e e e f f f f f f f f f f f f f f f f f f f f f f f
+// <------------------------------float-------------------------->
+// 3 3             2 2             1 1                           0
+// 1 0             3 2             5 4                           0
+//
+//
+// This type only supports conversion back and forth with float.
+//
+// This file must be compilable by nvcc.
+//
+// The type is defined in framework/numeric_types.h.
+
+namespace tensorflow {
+
+// Conversion routines between an array of float and bfloat16 of
+// "size".
+void FloatToBFloat16(const float* src, bfloat16* dst, int64 size);
+void BFloat16ToFloat(const bfloat16* src, float* dst, int64 size);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_BFLOAT16_H_

bfloat16_test.cc

@@ -0,0 +1,158 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/bfloat16.h"
+
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/lib/core/casts.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(Bfloat16Test, Simple) {
+  bfloat16 a(12);
+  // Floating point representation of 12: 0x41400000
+  EXPECT_EQ(0x4140, a.value);
+}
+
+float BinaryToFloat(uint32_t sign, uint32_t exponent, uint32_t high_mantissa,
+                    uint32_t low_mantissa) {
+  return bit_cast<float>((sign << 31) + (exponent << 23) +
+                         (high_mantissa << 16) + low_mantissa);
+}
+
+struct Bfloat16TestParam {
+  float input;
+  float expected;
+};
+
+class Bfloat16Test : public ::testing::Test,
+                     public ::testing::WithParamInterface<Bfloat16TestParam> {};
+
+TEST_P(Bfloat16Test, TruncateTest) {
+  bfloat16 a(GetParam().input);
+  if (std::isnan(GetParam().input)) {
+    EXPECT_TRUE(std::isnan(float(a)) || std::isinf(float(a)));
+    return;
+  }
+  EXPECT_EQ(GetParam().expected, float(a));
+}
+
+INSTANTIATE_TEST_CASE_P(
+    Bfloat16Test_Instantiation, Bfloat16Test,
+    ::testing::Values(
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b1111010111000011),
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(1, 0b10000000, 0b1001000, 0b1111010111000011),
+            BinaryToFloat(1, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b1000000000000000),
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b11111111, 0b0000000, 0b0000000000000001),
+            BinaryToFloat(0, 0b11111111, 0b0000000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b11111111, 0b1111111, 0b1111111111111111),
+            BinaryToFloat(0, 0b11111111, 0b1111111, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(1, 0b10000000, 0b1001000, 0b1100000000000000),
+            BinaryToFloat(1, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000),
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0100000000000000),
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b1000000000000000),
+            BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b00000000, 0b1001000, 0b1000000000000000),
+            BinaryToFloat(0, 0b00000000, 0b1001000, 0b0000000000000000)},
+        Bfloat16TestParam{
+            BinaryToFloat(0, 0b00000000, 0b1111111, 0b1100000000000000),
+            BinaryToFloat(0, 0b00000000, 0b1111111, 0b0000000000000000)}));
+
+TEST(Bfloat16Test, Conversion) {
+  float a[100];
+  for (int i = 0; i < 100; ++i) {
+    a[i] = i + 1.25;
+  }
+  bfloat16 b[100];
+  float c[100];
+  FloatToBFloat16(a, b, 100);
+  BFloat16ToFloat(b, c, 100);
+  for (int i = 0; i < 100; ++i) {
+    // The relative error should be less than 1/(2^7) since bfloat16
+    // has 7 bits mantissa.
+    EXPECT_LE(fabs(c[i] - a[i]) / a[i], 1.0 / 128);
+  }
+}
+
+TEST(Bfloat16Test, Epsilon) {
+  EXPECT_LT(1.0f, static_cast<float>(bfloat16::epsilon() + bfloat16(1.0f)));
+  EXPECT_EQ(1.0f, static_cast<float>((bfloat16::epsilon() / bfloat16(2.0f)) +
+                                     bfloat16(1.0f)));
+}
+
+TEST(Bfloat16Test, Negate) {
+  EXPECT_EQ(-3.0f, static_cast<float>(-bfloat16(3.0f)));
+  EXPECT_EQ(4.5f, static_cast<float>(-bfloat16(-4.5f)));
+}
+
+static void BM_FloatToBFloat16(int iters) {
+  testing::StopTiming();
+  static const int N = 32 << 20;
+  const int64 tot = static_cast<int64>(iters) * N;
+  testing::ItemsProcessed(tot);
+  testing::BytesProcessed(tot * (sizeof(float) + sizeof(bfloat16)));
+
+  float* inp = new float[N];
+  bfloat16* out = new bfloat16[N];
+
+  testing::StartTiming();
+  while (iters--) {
+    FloatToBFloat16(inp, out, N);
+  }
+  delete[] inp;
+  delete[] out;
+}
+BENCHMARK(BM_FloatToBFloat16);
+
+static void BM_BFloat16ToFloat(int iters) {
+  testing::StopTiming();
+  static const int N = 32 << 20;
+  const int64 tot = static_cast<int64>(iters) * N;
+  testing::ItemsProcessed(tot);
+  testing::BytesProcessed(tot * (sizeof(float) + sizeof(bfloat16)));
+
+  bfloat16* inp = new bfloat16[N];
+  float* out = new float[N];
+
+  testing::StartTiming();
+  while (iters--) {
+    BFloat16ToFloat(inp, out, N);
+  }
+  delete[] inp;
+  delete[] out;
+}
+BENCHMARK(BM_BFloat16ToFloat);
+
+}  // namespace
+}  // namespace tensorflow

cancellation.cc

@@ -0,0 +1,94 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/cancellation.h"
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+const CancellationToken CancellationManager::kInvalidToken = -1;
+
+CancellationManager::CancellationManager()
+    : is_cancelling_(false),
+      is_cancelled_(false),
+      next_cancellation_token_(0) {}
+
+void CancellationManager::StartCancel() {
+  gtl::FlatMap<CancellationToken, CancelCallback> callbacks_to_run;
+  {
+    mutex_lock l(mu_);
+    if (is_cancelled_.load(std::memory_order_relaxed) || is_cancelling_) {
+      return;
+    }
+    is_cancelling_ = true;
+    std::swap(callbacks_, callbacks_to_run);
+  }
+  // We call these callbacks without holding mu_, so that concurrent
+  // calls to DeregisterCallback, which can happen asynchronously, do
+  // not block. The callbacks remain valid because any concurrent call
+  // to DeregisterCallback will block until the
+  // cancelled_notification_ is notified.
+  for (auto key_and_value : callbacks_to_run) {
+    key_and_value.second();
+  }
+  {
+    mutex_lock l(mu_);
+    is_cancelling_ = false;
+    is_cancelled_.store(true, std::memory_order_release);
+  }
+  cancelled_notification_.Notify();
+}
+
+CancellationToken CancellationManager::get_cancellation_token() {
+  mutex_lock l(mu_);
+  return next_cancellation_token_++;
+}
+
+bool CancellationManager::RegisterCallback(CancellationToken token,
+                                           CancelCallback callback) {
+  mutex_lock l(mu_);
+  CHECK_LT(token, next_cancellation_token_) << "Invalid cancellation token";
+  bool should_register = !is_cancelled_ && !is_cancelling_;
+  if (should_register) {
+    std::swap(callbacks_[token], callback);
+  }
+  return should_register;
+}
+
+bool CancellationManager::DeregisterCallback(CancellationToken token) {
+  mu_.lock();
+  if (is_cancelled_) {
+    mu_.unlock();
+    return false;
+  } else if (is_cancelling_) {
+    mu_.unlock();
+    // Wait for all of the cancellation callbacks to be called. This
+    // wait ensures that the caller of DeregisterCallback does not
+    // return immediately and free objects that may be used in the
+    // execution of any currently pending callbacks in StartCancel.
+    cancelled_notification_.WaitForNotification();
+    return false;
+  } else {
+    callbacks_.erase(token);
+    mu_.unlock();
+    return true;
+  }
+}
+
+CancellationManager::~CancellationManager() { StartCancel(); }
+
+}  // end namespace tensorflow

cancellation.h

@@ -0,0 +1,137 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_CANCELLATION_H_
+#define TENSORFLOW_FRAMEWORK_CANCELLATION_H_
+
+#include <atomic>
+#include <functional>
+
+#include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// A token that can be used to register and deregister a
+// CancelCallback with a CancellationManager.
+//
+// CancellationToken values must be created by a call to
+// CancellationManager::get_cancellation_token.
+typedef int64 CancellationToken;
+
+// A callback that is invoked when a step is canceled.
+//
+// NOTE(mrry): See caveats about CancelCallback implementations in the
+// comment for CancellationManager::RegisterCallback.
+typedef std::function<void()> CancelCallback;
+
+class CancellationManager {
+ public:
+  // A value that won't be returned by get_cancellation_token().
+  static const CancellationToken kInvalidToken;
+
+  CancellationManager();
+  ~CancellationManager();
+
+  // Run all callbacks associated with this manager.
+  void StartCancel();
+
+  // Returns true iff StartCancel() has been called.
+  bool IsCancelled() { return is_cancelled_.load(std::memory_order_acquire); }
+
+  // Returns a token that must be used in calls to RegisterCallback
+  // and DeregisterCallback.
+  CancellationToken get_cancellation_token();
+
+  // Attempts to register the given callback to be invoked when this
+  // manager is cancelled. Returns true if the callback was
+  // registered; returns false if this manager was already cancelled,
+  // and the callback was not registered.
+  //
+  // If this method returns false, it is the caller's responsibility
+  // to perform any cancellation cleanup.
+  //
+  // This method is tricky to use correctly. The following usage pattern
+  // is recommended:
+  //
+  // class ObjectWithCancellableOperation {
+  //   mutex mu_;
+  //   void CancellableOperation(CancellationManager* cm,
+  //                             std::function<void(Status)> callback) {
+  //     bool already_cancelled;
+  //     CancellationToken token = cm->get_cancellation_token();
+  //     {
+  //       mutex_lock(mu_);
+  //       already_cancelled = !cm->RegisterCallback(
+  //           [this, token]() { Cancel(token); });
+  //       if (!already_cancelled) {
+  //         // Issue asynchronous operation. Associate the pending operation
+  //         // with `token` in some object state, or provide another way for
+  //         // the Cancel method to look up the operation for cancellation.
+  //         // Ensure that `cm->DeregisterCallback(token)` is called without
+  //         // holding `mu_`, before `callback` is invoked.
+  //         // ...
+  //       }
+  //     }
+  //     if (already_cancelled) {
+  //       callback(errors::Cancelled("Operation was cancelled"));
+  //     }
+  //   }
+  //
+  //   void Cancel(CancellationToken token) {
+  //     mutex_lock(mu_);
+  //     // Take action to cancel the operation with the given cancellation
+  //     // token.
+  //   }
+  //
+  // NOTE(mrry): The caller should take care that (i) the calling code
+  // is robust to `callback` being invoked asynchronously (e.g. from
+  // another thread), (ii) `callback` is deregistered by a call to
+  // this->DeregisterCallback(token) when the operation completes
+  // successfully, and (iii) `callback` does not invoke any method
+  // on this cancellation manager. Furthermore, it is important that
+  // the eventual caller of the complementary DeregisterCallback does not
+  // hold any mutexes that are required by `callback`.
+  bool RegisterCallback(CancellationToken token, CancelCallback callback);
+
+  // Deregister the callback that, when registered, was associated
+  // with the given cancellation token. Returns true iff the callback
+  // was deregistered and will not be invoked; otherwise returns false
+  // after the callback has been invoked, blocking if necessary.
+  //
+  // NOTE(mrry): This method may block if cancellation is in progress.
+  // The caller of this method must not hold any mutexes that are required
+  // to invoke any cancellation callback that has been registered with this
+  // cancellation manager.
+  bool DeregisterCallback(CancellationToken token);
+
+ private:
+  bool is_cancelling_;
+  std::atomic_bool is_cancelled_;
+
+  mutex mu_;
+  Notification cancelled_notification_;
+  CancellationToken next_cancellation_token_ GUARDED_BY(mu_);
+  gtl::FlatMap<CancellationToken, CancelCallback> callbacks_ GUARDED_BY(mu_);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_CANCELLATION_H_

cancellation_test.cc

@@ -0,0 +1,118 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/cancellation.h"
+
+#include <vector>
+#include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+TEST(Cancellation, SimpleNoCancel) {
+  bool is_cancelled = false;
+  CancellationManager* manager = new CancellationManager();
+  auto token = manager->get_cancellation_token();
+  bool registered = manager->RegisterCallback(
+      token, [&is_cancelled]() { is_cancelled = true; });
+  EXPECT_TRUE(registered);
+  bool deregistered = manager->DeregisterCallback(token);
+  EXPECT_TRUE(deregistered);
+  delete manager;
+  EXPECT_FALSE(is_cancelled);
+}
+
+TEST(Cancellation, SimpleCancel) {
+  bool is_cancelled = false;
+  CancellationManager* manager = new CancellationManager();
+  auto token = manager->get_cancellation_token();
+  bool registered = manager->RegisterCallback(
+      token, [&is_cancelled]() { is_cancelled = true; });
+  EXPECT_TRUE(registered);
+  manager->StartCancel();
+  EXPECT_TRUE(is_cancelled);
+  delete manager;
+}
+
+TEST(Cancellation, CancelBeforeRegister) {
+  CancellationManager* manager = new CancellationManager();
+  auto token = manager->get_cancellation_token();
+  manager->StartCancel();
+  bool registered = manager->RegisterCallback(token, nullptr);
+  EXPECT_FALSE(registered);
+  delete manager;
+}
+
+TEST(Cancellation, DeregisterAfterCancel) {
+  bool is_cancelled = false;
+  CancellationManager* manager = new CancellationManager();
+  auto token = manager->get_cancellation_token();
+  bool registered = manager->RegisterCallback(
+      token, [&is_cancelled]() { is_cancelled = true; });
+  EXPECT_TRUE(registered);
+  manager->StartCancel();
+  EXPECT_TRUE(is_cancelled);
+  bool deregistered = manager->DeregisterCallback(token);
+  EXPECT_FALSE(deregistered);
+  delete manager;
+}
+
+TEST(Cancellation, CancelMultiple) {
+  bool is_cancelled_1 = false, is_cancelled_2 = false, is_cancelled_3 = false;
+  CancellationManager* manager = new CancellationManager();
+  auto token_1 = manager->get_cancellation_token();
+  bool registered_1 = manager->RegisterCallback(
+      token_1, [&is_cancelled_1]() { is_cancelled_1 = true; });
+  EXPECT_TRUE(registered_1);
+  auto token_2 = manager->get_cancellation_token();
+  bool registered_2 = manager->RegisterCallback(
+      token_2, [&is_cancelled_2]() { is_cancelled_2 = true; });
+  EXPECT_TRUE(registered_2);
+  EXPECT_FALSE(is_cancelled_1);
+  EXPECT_FALSE(is_cancelled_2);
+  manager->StartCancel();
+  EXPECT_TRUE(is_cancelled_1);
+  EXPECT_TRUE(is_cancelled_2);
+  EXPECT_FALSE(is_cancelled_3);
+  auto token_3 = manager->get_cancellation_token();
+  bool registered_3 = manager->RegisterCallback(
+      token_3, [&is_cancelled_3]() { is_cancelled_3 = true; });
+  EXPECT_FALSE(registered_3);
+  EXPECT_FALSE(is_cancelled_3);
+  delete manager;
+}
+
+TEST(Cancellation, IsCancelled) {
+  CancellationManager* cm = new CancellationManager();
+  thread::ThreadPool w(Env::Default(), "test", 4);
+  std::vector<Notification> done(8);
+  for (size_t i = 0; i < done.size(); ++i) {
+    Notification* n = &done[i];
+    w.Schedule([n, cm]() {
+      while (!cm->IsCancelled()) {
+      }
+      n->Notify();
+    });
+  }
+  Env::Default()->SleepForMicroseconds(1000000 /* 1 second */);
+  cm->StartCancel();
+  for (size_t i = 0; i < done.size(); ++i) {
+    done[i].WaitForNotification();
+  }
+  delete cm;
+}
+
+}  // namespace tensorflow

common_shape_fns.cc

@@ -0,0 +1,1399 @@
+/* Copyright 2016 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.
+==============================================================================*/
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+
+namespace tensorflow {
+
+Status GetWindowedOutputSizeVerboseV2(int64 input_size, int64 filter_size,
+                                      int64 dilation_rate, int64 stride,
+                                      Padding padding_type, int64* output_size,
+                                      int64* padding_before,
+                                      int64* padding_after) {
+  if (stride <= 0) {
+    return errors::InvalidArgument("Stride must be > 0, but got ", stride);
+  }
+  if (dilation_rate < 1) {
+    return errors::InvalidArgument("Dilation rate must be >= 1, but got ",
+                                   dilation_rate);
+  }
+
+  // See also the parallel implementation in GetWindowedOutputSizeFromDimsV2.
+  int64 effective_filter_size = (filter_size - 1) * dilation_rate + 1;
+  switch (padding_type) {
+    case Padding::VALID:
+      *output_size = (input_size - effective_filter_size + stride) / stride;
+      *padding_before = *padding_after = 0;
+      break;
+    case Padding::SAME:
+      *output_size = (input_size + stride - 1) / stride;
+      const int64 padding_needed =
+          std::max(0LL, (*output_size - 1) * stride + effective_filter_size -
+                            input_size);
+      // For odd values of total padding, add more padding at the 'right'
+      // side of the given dimension.
+      *padding_before = padding_needed / 2;
+      *padding_after = padding_needed - *padding_before;
+      break;
+  }
+  if (*output_size < 0) {
+    return errors::InvalidArgument("computed output size would be negative");
+  }
+  return Status::OK();
+}
+
+Status GetWindowedOutputSizeVerbose(int64 input_size, int64 filter_size,
+                                    int64 stride, Padding padding_type,
+                                    int64* output_size, int64* padding_before,
+                                    int64* padding_after) {
+  return GetWindowedOutputSizeVerboseV2(input_size, filter_size,
+                                        /*dilation_rate=*/1, stride,
+                                        padding_type, output_size,
+                                        padding_before, padding_after);
+}
+
+Status GetWindowedOutputSize(int64 input_size, int64 filter_size, int64 stride,
+                             Padding padding_type, int64* output_size,
+                             int64* padding_size) {
+  int64 padding_after_unused;
+  return GetWindowedOutputSizeVerbose(input_size, filter_size, stride,
+                                      padding_type, output_size, padding_size,
+                                      &padding_after_unused);
+}
+
+Status GetWindowedOutputSizeV2(int64 input_size, int64 filter_size,
+                               int64 dilation_rate, int64 stride,
+                               Padding padding_type, int64* output_size,
+                               int64* padding_size) {
+  int64 padding_after_unused;
+  return GetWindowedOutputSizeVerboseV2(input_size, filter_size, dilation_rate,
+                                        stride, padding_type, output_size,
+                                        padding_size, &padding_after_unused);
+}
+
+Status Get3dOutputSize(const std::array<int64, 3>& input,
+                       const std::array<int64, 3>& window,
+                       const std::array<int64, 3>& strides,
+                       Padding padding_type, std::array<int64, 3>* output_ptr,
+                       std::array<int64, 3>* padding_ptr) {
+  for (size_t i = 0; i < input.size(); ++i) {
+    TF_RETURN_IF_ERROR(GetWindowedOutputSize(input[i], window[i], strides[i],
+                                             padding_type, &(*output_ptr)[i],
+                                             &(*padding_ptr)[i]));
+  }
+  return Status::OK();
+}
+
+Status Get3dOutputSizeV2(const std::array<int64, 3>& input,
+                         const std::array<int64, 3>& window,
+                         const std::array<int64, 3>& dilations,
+                         const std::array<int64, 3>& strides,
+                         Padding padding_type, std::array<int64, 3>* output_ptr,
+                         std::array<int64, 3>* padding_ptr) {
+  for (size_t i = 0; i < input.size(); ++i) {
+    TF_RETURN_IF_ERROR(GetWindowedOutputSizeV2(
+        input[i], window[i], dilations[i], strides[i], padding_type,
+        &(*output_ptr)[i], &(*padding_ptr)[i]));
+  }
+  return Status::OK();
+}
+
+namespace shape_inference {
+
+// The V2 version computes windowed output size with arbitrary dilation_rate,
+// while the original version only handles the cases where dilation_rates equal
+// to 1.
+Status GetWindowedOutputSizeFromDimsV2(
+    shape_inference::InferenceContext* c,
+    shape_inference::DimensionHandle input_size,
+    shape_inference::DimensionOrConstant filter_size, int64 dilation_rate,
+    int64 stride, Padding padding_type,
+    shape_inference::DimensionHandle* output_size) {
+  if (stride <= 0) {
+    return errors::InvalidArgument("Stride must be > 0, but got ", stride);
+  }
+
+  if (dilation_rate < 1) {
+    return errors::InvalidArgument("Dilation rate must be >= 1, but got ",
+                                   dilation_rate);
+  }
+
+  // See also the parallel implementation in GetWindowedOutputSizeVerbose.
+  switch (padding_type) {
+    case Padding::VALID:
+      if (dilation_rate > 1) {
+        DimensionHandle window_size;
+        TF_RETURN_IF_ERROR(
+            c->Subtract(c->MakeDim(filter_size), 1, &window_size));
+        TF_RETURN_IF_ERROR(
+            c->Multiply(window_size, dilation_rate, &window_size));
+        TF_RETURN_IF_ERROR(c->Add(window_size, 1, &window_size));
+        TF_RETURN_IF_ERROR(c->Subtract(input_size, window_size, output_size));
+      } else {
+        TF_RETURN_IF_ERROR(c->Subtract(input_size, filter_size, output_size));
+      }
+      TF_RETURN_IF_ERROR(c->Add(*output_size, stride, output_size));
+      TF_RETURN_IF_ERROR(c->Divide(*output_size, stride,
+                                   /*evenly_divisible=*/false, output_size));
+      break;
+    case Padding::SAME:
+      TF_RETURN_IF_ERROR(c->Add(input_size, stride - 1, output_size));
+      TF_RETURN_IF_ERROR(c->Divide(*output_size, stride,
+                                   /*evenly_divisible=*/false, output_size));
+      break;
+  }
+  return Status::OK();
+}
+
+Status GetWindowedOutputSizeFromDims(
+    shape_inference::InferenceContext* c,
+    shape_inference::DimensionHandle input_size,
+    shape_inference::DimensionOrConstant filter_size, int64 stride,
+    Padding padding_type, shape_inference::DimensionHandle* output_size) {
+  return GetWindowedOutputSizeFromDimsV2(c, input_size, filter_size,
+                                         /*dilation_rate=*/1, stride,
+                                         padding_type, output_size);
+}
+
+Status UnchangedShape(shape_inference::InferenceContext* c) {
+  c->set_output(0, c->input(0));
+  return Status::OK();
+}
+
+Status MatMulShape(shape_inference::InferenceContext* c) {
+  ShapeHandle a;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a));
+
+  ShapeHandle b;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b));
+
+  bool transpose_a, transpose_b;
+  TF_RETURN_IF_ERROR(c->GetAttr("transpose_a", &transpose_a));
+  TF_RETURN_IF_ERROR(c->GetAttr("transpose_b", &transpose_b));
+  DimensionHandle output_rows = transpose_a ? c->Dim(a, 1) : c->Dim(a, 0);
+  DimensionHandle output_cols = transpose_b ? c->Dim(b, 0) : c->Dim(b, 1);
+
+  // Validate that the inner shapes are compatible.
+  DimensionHandle inner_a = transpose_a ? c->Dim(a, 0) : c->Dim(a, 1);
+  DimensionHandle inner_b = transpose_b ? c->Dim(b, 1) : c->Dim(b, 0);
+  DimensionHandle merged;
+  TF_RETURN_IF_ERROR(c->Merge(inner_a, inner_b, &merged));
+
+  c->set_output(0, c->Matrix(output_rows, output_cols));
+  return Status::OK();
+}
+
+Status BiasAddShape(shape_inference::InferenceContext* c) {
+  ShapeHandle input_shape;
+
+  // Fetch the data_format attribute, which may not exist.
+  string data_format;
+  Status s = c->GetAttr("data_format", &data_format);
+
+  if (s.ok() && data_format == "NCHW") {
+    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape));
+  } else {
+    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
+  }
+
+  ShapeHandle bias_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &bias_shape));
+  DimensionHandle bias_dim = c->Dim(bias_shape, 0);
+
+  // If rank unknown, return unknown shape.
+  if (!c->RankKnown(input_shape)) {
+    c->set_output(0, c->UnknownShape());
+    return Status::OK();
+  }
+
+  // Output has the same shape as the input, and matches the length of
+  // the bias in its bias dimension.
+  ShapeHandle output_shape;
+  if (s.ok() && data_format == "NCHW") {
+    // Merge the length of bias_shape into the third to last dimension
+    ShapeHandle first;
+    TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, -3, &first));
+
+    ShapeHandle last;
+    TF_RETURN_IF_ERROR(c->Subshape(input_shape, -2, &last));
+
+    DimensionHandle input_bias_dim = c->Dim(input_shape, -3);
+    DimensionHandle merged_bias_dim;
+    TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim));
+    ShapeHandle merged_bias = c->Vector(merged_bias_dim);
+
+    ShapeHandle temp;
+    TF_RETURN_IF_ERROR(c->Concatenate(first, merged_bias, &temp));
+    TF_RETURN_IF_ERROR(c->Concatenate(temp, last, &output_shape));
+  } else {
+    ShapeHandle all_but_bias;
+    TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, -1, &all_but_bias));
+
+    DimensionHandle input_bias_dim = c->Dim(input_shape, -1);
+    DimensionHandle merged_bias_dim;
+    TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim));
+
+    ShapeHandle merged_bias = c->Vector(merged_bias_dim);
+    TF_RETURN_IF_ERROR(
+        c->Concatenate(all_but_bias, merged_bias, &output_shape));
+  }
+
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status BiasAddGradShape(shape_inference::InferenceContext* c) {
+  ShapeHandle input_shape;
+  // Fetch the data_format attribute, which may not exist.
+  string data_format;
+  Status s = c->GetAttr("data_format", &data_format);
+
+  if (s.ok() && data_format == "NCHW") {
+    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape));
+    c->set_output(0, c->Vector(c->Dim(input_shape, -3)));
+  } else {
+    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
+    c->set_output(0, c->Vector(c->Dim(input_shape, -1)));
+  }
+
+  return Status::OK();
+}
+
+Status CheckFormatConstraintsOnShape(const TensorFormat tensor_format,
+                                     const ShapeHandle shape_handle,
+                                     const string& tensor_name,
+                                     shape_inference::InferenceContext* c) {
+  if (tensor_format == FORMAT_NCHW_VECT_C) {
+    // Check that the vect dim has size 4.
+    const int num_dims = c->Rank(shape_handle);
+    DimensionHandle vect_dim = c->Dim(
+        shape_handle, GetTensorInnerFeatureDimIndex(num_dims, tensor_format));
+    DimensionHandle unused_vect_dim;
+    TF_RETURN_IF_ERROR(c->WithValue(vect_dim, 4, &unused_vect_dim));
+  }
+
+  return Status::OK();
+}
+
+Status MakeShapeFromFormat(TensorFormat format, DimensionOrConstant N,
+                           const std::vector<DimensionOrConstant>& spatial,
+                           DimensionOrConstant C, ShapeHandle* out,
+                           shape_inference::InferenceContext* context) {
+  const int num_dims = GetTensorDimsFromSpatialDims(spatial.size(), format);
+  std::vector<DimensionHandle> dims_actual(num_dims);
+  dims_actual[GetTensorBatchDimIndex(num_dims, format)] = context->MakeDim(N);
+  int outer_c_index = GetTensorFeatureDimIndex(num_dims, format);
+  dims_actual[outer_c_index] = context->MakeDim(C);
+  if (format == FORMAT_NCHW_VECT_C) {
+    dims_actual[GetTensorInnerFeatureDimIndex(num_dims, format)] =
+        context->MakeDim(4);
+  }
+  for (int spatial_dim = 0; spatial_dim < spatial.size(); spatial_dim++) {
+    dims_actual[GetTensorSpatialDimIndex(num_dims, format, spatial_dim)] =
+        context->MakeDim(spatial[spatial_dim]);
+  }
+  *out = context->MakeShape(dims_actual);
+  return Status::OK();
+}
+
+Status DimensionsFromShape(ShapeHandle shape, TensorFormat format,
+                           DimensionHandle* batch_dim,
+                           gtl::MutableArraySlice<DimensionHandle> spatial_dims,
+                           DimensionHandle* filter_dim,
+                           InferenceContext* context) {
+  const int32 rank = GetTensorDimsFromSpatialDims(spatial_dims.size(), format);
+  // Batch.
+  *batch_dim = context->Dim(shape, GetTensorBatchDimIndex(rank, format));
+  // Spatial.
+  for (int spatial_dim_index = 0; spatial_dim_index < spatial_dims.size();
+       ++spatial_dim_index) {
+    spatial_dims[spatial_dim_index] = context->Dim(
+        shape, GetTensorSpatialDimIndex(rank, format, spatial_dim_index));
+  }
+  // Channel.
+  *filter_dim = context->Dim(shape, GetTensorFeatureDimIndex(rank, format));
+  if (format == FORMAT_NCHW_VECT_C) {
+    TF_RETURN_IF_ERROR(context->Multiply(
+        *filter_dim,
+        context->Dim(shape, GetTensorInnerFeatureDimIndex(rank, format)),
+        filter_dim));
+  }
+  return Status::OK();
+}
+
+Status ShapeFromDimensions(DimensionHandle batch_dim,
+                           gtl::ArraySlice<DimensionHandle> spatial_dims,
+                           DimensionHandle filter_dim, TensorFormat format,
+                           InferenceContext* context, ShapeHandle* shape) {
+  const int32 rank = GetTensorDimsFromSpatialDims(spatial_dims.size(), format);
+  std::vector<DimensionHandle> out_dims(rank);
+
+  // Batch.
+  out_dims[tensorflow::GetTensorBatchDimIndex(rank, format)] = batch_dim;
+  // Spatial.
+  for (int spatial_dim_index = 0; spatial_dim_index < spatial_dims.size();
+       ++spatial_dim_index) {
+    out_dims[tensorflow::GetTensorSpatialDimIndex(
+        rank, format, spatial_dim_index)] = spatial_dims[spatial_dim_index];
+  }
+  // Channel.
+  if (format == tensorflow::FORMAT_NCHW_VECT_C) {
+    // When format is NCHW_VECT_C, factor the feature map count
+    // into the outer feature count and the inner feature count (=4).
+    TF_RETURN_IF_ERROR(context->Divide(
+        filter_dim, 4, /*evenly_divisible=*/true,
+        &out_dims[tensorflow::GetTensorFeatureDimIndex(rank, format)]));
+    out_dims[GetTensorInnerFeatureDimIndex(rank, format)] = context->MakeDim(4);
+  } else {
+    out_dims[tensorflow::GetTensorFeatureDimIndex(rank, format)] = filter_dim;
+  }
+
+  *shape = context->MakeShape(out_dims);
+  return tensorflow::Status::OK();
+}
+
+Status Conv2DShape(shape_inference::InferenceContext* c) {
+  string data_format_str, filter_format_str;
+  if (!c->GetAttr("data_format", &data_format_str).ok()) {
+    data_format_str = "NHWC";
+  }
+  if (!c->GetAttr("filter_format", &filter_format_str).ok()) {
+    filter_format_str = "HWIO";
+  }
+
+  TensorFormat data_format;
+  if (!FormatFromString(data_format_str, &data_format)) {
+    return errors::InvalidArgument("Invalid data format string: ",
+                                   data_format_str);
+  }
+  FilterTensorFormat filter_format;
+  if (!FilterFormatFromString(filter_format_str, &filter_format)) {
+    return errors::InvalidArgument("Invalid filter format string: ",
+                                   filter_format_str);
+  }
+
+  constexpr int num_spatial_dims = 2;
+  const int rank = GetTensorDimsFromSpatialDims(num_spatial_dims, data_format);
+  ShapeHandle conv_input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &conv_input_shape));
+  TF_RETURN_IF_ERROR(CheckFormatConstraintsOnShape(
+      data_format, conv_input_shape, "conv_input", c));
+
+  // The filter rank should match the input (4 for NCHW, 5 for NCHW_VECT_C).
+  ShapeHandle filter_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), rank, &filter_shape));
+  TF_RETURN_IF_ERROR(
+      CheckFormatConstraintsOnShape(data_format, filter_shape, "filter", c));
+
+  std::vector<int32> dilations;
+  TF_RETURN_IF_ERROR(c->GetAttr("dilations", &dilations));
+
+  if (dilations.size() != 4) {
+    return errors::InvalidArgument(
+        "Conv2D requires the dilation attribute to contain 4 values, but got: ",
+        dilations.size());
+  }
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+
+  // strides.size() should be 4 (NCHW) even if the input is 5 (NCHW_VECT_C).
+  if (strides.size() != 4) {
+    return errors::InvalidArgument("Conv2D on data format ", data_format_str,
+                                   " requires the stride attribute to contain"
+                                   " 4 values, but got: ",
+                                   strides.size());
+  }
+
+  const int32 stride_rows = GetTensorDim(strides, data_format, 'H');
+  const int32 stride_cols = GetTensorDim(strides, data_format, 'W');
+  const int32 dilation_rows = GetTensorDim(dilations, data_format, 'H');
+  const int32 dilation_cols = GetTensorDim(dilations, data_format, 'W');
+
+  DimensionHandle batch_size_dim;
+  DimensionHandle input_depth_dim;
+  gtl::InlinedVector<DimensionHandle, 2> input_spatial_dims(2);
+  TF_RETURN_IF_ERROR(DimensionsFromShape(conv_input_shape, data_format,
+                                         &batch_size_dim, &input_spatial_dims,
+                                         &input_depth_dim, c));
+
+  DimensionHandle output_depth_dim = c->Dim(
+      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'O'));
+  DimensionHandle filter_rows_dim = c->Dim(
+      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'H'));
+  DimensionHandle filter_cols_dim = c->Dim(
+      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'W'));
+  DimensionHandle filter_input_depth_dim;
+  if (filter_format == FORMAT_OIHW_VECT_I) {
+    TF_RETURN_IF_ERROR(c->Multiply(
+        c->Dim(filter_shape,
+               GetFilterDimIndex<num_spatial_dims>(filter_format, 'I')),
+        c->Dim(filter_shape,
+               GetFilterTensorInnerInputChannelsDimIndex(rank, filter_format)),
+        &filter_input_depth_dim));
+  } else {
+    filter_input_depth_dim = c->Dim(
+        filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'I'));
+  }
+
+  // Check that the input tensor and the filter tensor agree on the input
+  // channel count.
+  DimensionHandle unused;
+  TF_RETURN_IF_ERROR(
+      c->Merge(input_depth_dim, filter_input_depth_dim, &unused));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  DimensionHandle output_rows, output_cols;
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
+      c, input_spatial_dims[0], filter_rows_dim, dilation_rows, stride_rows,
+      padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
+      c, input_spatial_dims[1], filter_cols_dim, dilation_cols, stride_cols,
+      padding, &output_cols));
+
+  ShapeHandle output_shape;
+  TF_RETURN_IF_ERROR(
+      ShapeFromDimensions(batch_size_dim, {output_rows, output_cols},
+                          output_depth_dim, data_format, c, &output_shape));
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+// TODO(mjanusz): Unify all conv/pooling shape functions.
+Status Conv3DShape(shape_inference::InferenceContext* c) {
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input_shape));
+  ShapeHandle filter_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 5, &filter_shape));
+
+  string data_format;
+  Status s = c->GetAttr("data_format", &data_format);
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+  if (strides.size() != 5) {
+    return errors::InvalidArgument(
+        "Conv3D requires the stride attribute to contain 5 values, but got: ",
+        strides.size());
+  }
+
+  int32 stride_planes, stride_rows, stride_cols;
+  if (s.ok() && data_format == "NCDHW") {
+    // Convert input_shape to NDHWC.
+    auto dim = [&](char dimension) {
+      return c->Dim(input_shape, GetTensorDimIndex<3>(FORMAT_NCHW, dimension));
+    };
+    input_shape =
+        c->MakeShape({{dim('N'), dim('0'), dim('1'), dim('2'), dim('C')}});
+    stride_planes = strides[2];
+    stride_cols = strides[3];
+    stride_rows = strides[4];
+  } else {
+    stride_planes = strides[1];
+    stride_rows = strides[2];
+    stride_cols = strides[3];
+  }
+
+  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
+  DimensionHandle in_planes_dim = c->Dim(input_shape, 1);
+  DimensionHandle in_rows_dim = c->Dim(input_shape, 2);
+  DimensionHandle in_cols_dim = c->Dim(input_shape, 3);
+
+  DimensionHandle filter_planes_dim = c->Dim(filter_shape, 0);
+  DimensionHandle filter_rows_dim = c->Dim(filter_shape, 1);
+  DimensionHandle filter_cols_dim = c->Dim(filter_shape, 2);
+  DimensionHandle output_depth_dim = c->Dim(filter_shape, 4);
+
+  DimensionHandle unused;
+  TF_RETURN_IF_ERROR(
+      c->Merge(c->Dim(input_shape, 4), c->Dim(filter_shape, 3), &unused));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+  DimensionHandle output_planes, output_rows, output_cols;
+
+  TF_RETURN_IF_ERROR(
+      GetWindowedOutputSizeFromDims(c, in_planes_dim, filter_planes_dim,
+                                    stride_planes, padding, &output_planes));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, filter_rows_dim, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, filter_cols_dim, stride_cols, padding, &output_cols));
+
+  ShapeHandle output_shape;
+  if (data_format == "NCDHW") {
+    output_shape = c->MakeShape({batch_size_dim, output_depth_dim,
+                                 output_planes, output_rows, output_cols});
+  } else {
+    output_shape = c->MakeShape({batch_size_dim, output_planes, output_rows,
+                                 output_cols, output_depth_dim});
+  }
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status DepthwiseConv2DNativeShape(shape_inference::InferenceContext* c) {
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input_shape));
+  ShapeHandle filter_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 4, &filter_shape));
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+
+  if (strides.size() != 4) {
+    return errors::InvalidArgument(
+        "DepthwiseConv2D requires the stride attribute to contain 4 values, "
+        "but got: ",
+        strides.size());
+  }
+
+  string data_format;
+  Status s = c->GetAttr("data_format", &data_format);
+  int32 stride_rows;
+  int32 stride_cols;
+  if (s.ok() && data_format == "NCHW") {
+    // Canonicalize input shape to NHWC so the shape inference code below can
+    // process it.
+    input_shape =
+        c->MakeShape({{c->Dim(input_shape, 0), c->Dim(input_shape, 2),
+                       c->Dim(input_shape, 3), c->Dim(input_shape, 1)}});
+    stride_rows = strides[2];
+    stride_cols = strides[3];
+  } else {
+    stride_rows = strides[1];
+    stride_cols = strides[2];
+  }
+
+  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
+  DimensionHandle in_rows_dim = c->Dim(input_shape, 1);
+  DimensionHandle in_cols_dim = c->Dim(input_shape, 2);
+
+  DimensionHandle filter_rows_dim = c->Dim(filter_shape, 0);
+  DimensionHandle filter_cols_dim = c->Dim(filter_shape, 1);
+  DimensionHandle input_depth = c->Dim(filter_shape, 2);
+  DimensionHandle depth_multiplier = c->Dim(filter_shape, 3);
+
+  // Check that the input depths are compatible.
+  TF_RETURN_IF_ERROR(
+      c->Merge(c->Dim(input_shape, 3), input_depth, &input_depth));
+
+  DimensionHandle output_depth;
+  TF_RETURN_IF_ERROR(c->Multiply(input_depth, depth_multiplier, &output_depth));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  // TODO(mrry,shlens): Raise an error if the stride would cause
+  // information in the input to be ignored. This will require a change
+  // in the kernel implementation.
+  DimensionHandle output_rows, output_cols;
+
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, filter_rows_dim, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, filter_cols_dim, stride_cols, padding, &output_cols));
+
+  ShapeHandle output_shape;
+  if (data_format == "NCHW") {
+    output_shape =
+        c->MakeShape({batch_size_dim, output_depth, output_rows, output_cols});
+  } else {
+    output_shape =
+        c->MakeShape({batch_size_dim, output_rows, output_cols, output_depth});
+  }
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status AvgPoolShape(shape_inference::InferenceContext* c) {
+  string data_format_str;
+  TensorFormat data_format;
+  Status s = c->GetAttr("data_format", &data_format_str);
+  if (s.ok()) {
+    FormatFromString(data_format_str, &data_format);
+  } else {
+    data_format = FORMAT_NHWC;
+  }
+
+  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));
+
+  TF_RETURN_IF_ERROR(
+      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+  if (strides.size() != 4) {
+    return errors::InvalidArgument(
+        "AvgPool requires the stride attribute to contain 4 values, but got: ",
+        strides.size());
+  }
+
+  std::vector<int32> kernel_sizes;
+  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
+  if (kernel_sizes.size() != 4) {
+    return errors::InvalidArgument(
+        "AvgPool requires the ksize attribute to contain 4 values, but got: ",
+        kernel_sizes.size());
+  }
+
+  int32 stride_rows = GetTensorDim(strides, data_format, 'H');
+  int32 stride_cols = GetTensorDim(strides, data_format, 'W');
+  int32 kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
+  int32 kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');
+
+  constexpr int num_spatial_dims = 2;
+  DimensionHandle batch_size_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
+  DimensionHandle in_rows_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
+  DimensionHandle in_cols_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
+  DimensionHandle depth_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  // TODO(mrry,shlens): Raise an error if the stride would cause
+  // information in the input to be ignored. This will require a change
+  // in the kernel implementation.
+
+  DimensionHandle output_rows, output_cols;
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));
+
+  ShapeHandle output_shape;
+  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
+                                         {output_rows, output_cols}, depth_dim,
+                                         &output_shape, c));
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status FusedBatchNormShape(shape_inference::InferenceContext* c) {
+  ShapeHandle x;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &x));
+
+  bool is_training;
+  TF_RETURN_IF_ERROR(c->GetAttr("is_training", &is_training));
+  int number_inputs = (is_training) ? 3 : 5;
+  string data_format;
+  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format));
+  DimensionHandle channel_dim =
+      (data_format == "NHWC") ? c->Dim(x, 3) : c->Dim(x, 1);
+
+  // covers scale, offset, and if is_training is false, mean, variance
+  for (int i = 1; i < number_inputs; ++i) {
+    ShapeHandle vec;
+    TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &vec));
+    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(vec, 0), &channel_dim));
+  }
+
+  ShapeHandle y;
+  if (data_format == "NHWC") {
+    TF_RETURN_IF_ERROR(c->ReplaceDim(x, 3, channel_dim, &y));
+  } else {
+    TF_RETURN_IF_ERROR(c->ReplaceDim(x, 1, channel_dim, &y));
+  }
+  c->set_output(0, y);
+  ShapeHandle vector_shape = c->Vector(channel_dim);
+  c->set_output(1, vector_shape);
+  c->set_output(2, vector_shape);
+  c->set_output(3, vector_shape);
+  c->set_output(4, vector_shape);
+  return Status::OK();
+}
+
+Status FusedBatchNormGradShape(shape_inference::InferenceContext* c) {
+  ShapeHandle y_backprop;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &y_backprop));
+  ShapeHandle x;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 4, &x));
+
+  bool is_training;
+  string data_format;
+  TF_RETURN_IF_ERROR(c->GetAttr("is_training", &is_training));
+  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format));
+  DimensionHandle channel_dim =
+      (data_format == "NHWC") ? c->Dim(y_backprop, 3) : c->Dim(y_backprop, 1);
+  if (data_format == "NHWC") {
+    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(x, 3), &channel_dim));
+  } else {
+    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(x, 1), &channel_dim));
+  }
+
+  // covers scale, mean (reserve_space_1), variance (reserve_space_2)
+  for (int i = 2; i < 5; ++i) {
+    ShapeHandle vec;
+    TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &vec));
+    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(vec, 0), &channel_dim));
+  }
+
+  ShapeHandle x_backprop;
+  if (data_format == "NHWC") {
+    TF_RETURN_IF_ERROR(c->ReplaceDim(y_backprop, 3, channel_dim, &x_backprop));
+  } else {
+    TF_RETURN_IF_ERROR(c->ReplaceDim(y_backprop, 1, channel_dim, &x_backprop));
+  }
+  c->set_output(0, x_backprop);
+  c->set_output(1, c->Vector(channel_dim));
+  c->set_output(2, c->Vector(channel_dim));
+  // Set the correct shapes for reserve_spaces
+  // so that gradients can be performed when
+  // the op is in a symbolic condition.
+  if (is_training) {
+    c->set_output(3, c->Vector(0));
+    c->set_output(4, c->Vector(0));
+  } else {
+    c->set_output(3, c->Vector(channel_dim));
+    c->set_output(4, c->Vector(channel_dim));
+  }
+  return Status::OK();
+}
+
+Status MaxPoolShape(shape_inference::InferenceContext* c) {
+  string data_format_str;
+  TensorFormat data_format;
+  Status s = c->GetAttr("data_format", &data_format_str);
+  if (s.ok()) {
+    FormatFromString(data_format_str, &data_format);
+  } else {
+    data_format = FORMAT_NHWC;
+  }
+
+  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));
+
+  TF_RETURN_IF_ERROR(
+      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+  if (strides.size() != 4) {
+    return errors::InvalidArgument(
+        "MaxPool requires the stride attribute to contain 4 values, but got: ",
+        strides.size());
+  }
+
+  std::vector<int32> kernel_sizes;
+  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
+  if (kernel_sizes.size() != 4) {
+    return errors::InvalidArgument(
+        "MaxPool requires the ksize attribute to contain 4 values, but got: ",
+        kernel_sizes.size());
+  }
+
+  int32 stride_depth = GetTensorDim(strides, data_format, 'C');
+  int32 stride_rows = GetTensorDim(strides, data_format, 'H');
+  int32 stride_cols = GetTensorDim(strides, data_format, 'W');
+  int32 kernel_depth = GetTensorDim(kernel_sizes, data_format, 'C');
+  int32 kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
+  int32 kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');
+
+  constexpr int num_spatial_dims = 2;
+  DimensionHandle batch_size_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
+  DimensionHandle in_rows_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
+  DimensionHandle in_cols_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
+  DimensionHandle in_depth_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  ShapeHandle output_shape;
+  DimensionHandle output_rows, output_cols, output_depth;
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_depth_dim, kernel_depth, stride_depth, padding, &output_depth));
+
+  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
+                                         {output_rows, output_cols},
+                                         output_depth, &output_shape, c));
+
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status MaxPoolV2Shape(shape_inference::InferenceContext* c, int num_inputs) {
+  string data_format_str;
+  TensorFormat data_format;
+  Status s = c->GetAttr("data_format", &data_format_str);
+  if (s.ok()) {
+    FormatFromString(data_format_str, &data_format);
+  } else {
+    data_format = FORMAT_NHWC;
+  }
+
+  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));
+
+  TF_RETURN_IF_ERROR(
+      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));
+
+  std::vector<int32> kernel_sizes;
+  std::vector<int32> strides;
+
+  if (c->num_inputs() + 2 == num_inputs) {
+    TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
+
+    TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+  } else {
+    // Verify shape of ksize and strides input.
+    ShapeHandle size;
+    DimensionHandle unused;
+    TF_RETURN_IF_ERROR(c->WithRank(c->input(c->num_inputs() - 2), 1, &size));
+    TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 4, &unused));
+    TF_RETURN_IF_ERROR(c->WithRank(c->input(c->num_inputs() - 1), 1, &size));
+    TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 4, &unused));
+
+    const Tensor* kernel_sizes_tensor = c->input_tensor(c->num_inputs() - 2);
+    if (kernel_sizes_tensor == nullptr) {
+      c->set_output(0, c->UnknownShape());
+      return Status::OK();
+    }
+    kernel_sizes.resize(kernel_sizes_tensor->shape().num_elements());
+    auto kernel_sizes_vec = kernel_sizes_tensor->flat<int32>();
+    std::copy_n(&kernel_sizes_vec(0), kernel_sizes.size(),
+                kernel_sizes.begin());
+
+    const Tensor* strides_tensor = c->input_tensor(c->num_inputs() - 1);
+    if (strides_tensor == nullptr) {
+      c->set_output(0, c->UnknownShape());
+      return Status::OK();
+    }
+    strides.resize(strides_tensor->shape().num_elements());
+    auto strides_vec = strides_tensor->flat<int32>();
+    std::copy_n(&strides_vec(0), strides.size(), strides.begin());
+  }
+
+  if (strides.size() != 4) {
+    return errors::InvalidArgument(
+        "MaxPool requires the stride attribute to contain 4 values, but "
+        "got: ",
+        strides.size());
+  }
+  if (kernel_sizes.size() != 4) {
+    return errors::InvalidArgument(
+        "MaxPool requires the ksize attribute to contain 4 values, but got: ",
+        kernel_sizes.size());
+  }
+
+  int32 stride_depth = GetTensorDim(strides, data_format, 'C');
+  int32 stride_rows = GetTensorDim(strides, data_format, 'H');
+  int32 stride_cols = GetTensorDim(strides, data_format, 'W');
+  int32 kernel_depth = GetTensorDim(kernel_sizes, data_format, 'C');
+  int32 kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
+  int32 kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');
+
+  constexpr int num_spatial_dims = 2;
+  DimensionHandle batch_size_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
+  DimensionHandle in_rows_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
+  DimensionHandle in_cols_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
+  DimensionHandle in_depth_dim = c->Dim(
+      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  ShapeHandle output_shape;
+  DimensionHandle output_rows, output_cols, output_depth;
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_depth_dim, kernel_depth, stride_depth, padding, &output_depth));
+
+  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
+                                         {output_rows, output_cols},
+                                         output_depth, &output_shape, c));
+
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status Pool3DShape(shape_inference::InferenceContext* c) {
+  ShapeHandle input_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input_shape));
+
+  string data_format;
+  Status s = c->GetAttr("data_format", &data_format);
+
+  std::vector<int32> strides;
+  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
+  if (strides.size() != 5) {
+    return errors::InvalidArgument(
+        "Pool3D ops require the stride attribute to contain 5 values, but "
+        "got: ",
+        strides.size());
+  }
+
+  std::vector<int32> kernel_sizes;
+  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
+  if (kernel_sizes.size() != 5) {
+    return errors::InvalidArgument(
+        "Pool3D requires the ksize attribute to contain 5 values, but got: ",
+        kernel_sizes.size());
+  }
+
+  int32 stride_planes, stride_rows, stride_cols;
+  int32 kernel_planes, kernel_rows, kernel_cols;
+
+  if (s.ok() && data_format == "NCDHW") {
+    // Convert input_shape to NDHWC.
+    auto dim = [&](char dimension) {
+      return c->Dim(input_shape, GetTensorDimIndex<3>(FORMAT_NCHW, dimension));
+    };
+    input_shape =
+        c->MakeShape({{dim('N'), dim('0'), dim('1'), dim('2'), dim('C')}});
+    stride_planes = strides[2];
+    stride_rows = strides[3];
+    stride_cols = strides[4];
+    kernel_planes = kernel_sizes[2];
+    kernel_rows = kernel_sizes[3];
+    kernel_cols = kernel_sizes[4];
+  } else {
+    stride_planes = strides[1];
+    stride_rows = strides[2];
+    stride_cols = strides[3];
+    kernel_planes = kernel_sizes[1];
+    kernel_rows = kernel_sizes[2];
+    kernel_cols = kernel_sizes[3];
+  }
+
+  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
+  DimensionHandle in_planes_dim = c->Dim(input_shape, 1);
+  DimensionHandle in_rows_dim = c->Dim(input_shape, 2);
+  DimensionHandle in_cols_dim = c->Dim(input_shape, 3);
+  DimensionHandle output_depth_dim = c->Dim(input_shape, 4);
+
+  Padding padding;
+  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
+
+  // TODO(mrry,shlens): Raise an error if the stride would cause
+  // information in the input to be ignored. This will require a change
+  // in the kernel implementation.
+  DimensionHandle output_planes, output_rows, output_cols;
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_planes_dim, kernel_planes, stride_planes, padding, &output_planes));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
+  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
+      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));
+
+  ShapeHandle output_shape;
+  if (data_format == "NCDHW") {
+    output_shape = c->MakeShape({batch_size_dim, output_depth_dim,
+                                 output_planes, output_rows, output_cols});
+  } else {
+    output_shape = c->MakeShape({batch_size_dim, output_planes, output_rows,
+                                 output_cols, output_depth_dim});
+  }
+
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+Status UnknownShape(shape_inference::InferenceContext* c) {
+  for (int i = 0; i < c->num_outputs(); ++i) {
+    c->set_output(i, c->UnknownShape());
+  }
+  return Status::OK();
+}
+
+template <typename T>
+Status ReductionShapeHelper(const Tensor* reduction_indices_t,
+                            const int32 input_rank,
+                            std::set<int64>& true_indices) {
+  auto reduction_indices = reduction_indices_t->flat<T>();
+  for (int i = 0; i < reduction_indices_t->NumElements(); ++i) {
+    const T reduction_index = reduction_indices(i);
+    if (reduction_index < -input_rank || reduction_index >= input_rank) {
+      return errors::InvalidArgument("Invalid reduction dimension ",
+                                     reduction_index, " for input with ",
+                                     input_rank, " dimensions.");
+    }
+
+    auto wrapped_index = reduction_index;
+    if (wrapped_index < 0) {
+      wrapped_index += input_rank;
+    }
+
+    true_indices.insert(wrapped_index);
+  }
+  return Status::OK();
+}
+
+Status ReductionShape(InferenceContext* c) {
+  ShapeHandle input = c->input(0);
+
+  ShapeHandle indices;
+  // Older versions of TensorFlow accidentally allowed higher rank tensors like
+  // [[1,2]] or [[1],[2]] to represent axis=[1,2].
+  if (c->graph_def_version() < 21) {
+    indices = c->input(1);
+  } else {
+    TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &indices));
+  }
+
+  bool keep_dims;
+  TF_RETURN_IF_ERROR(c->GetAttr("keep_dims", &keep_dims));
+
+  const Tensor* reduction_indices_t = c->input_tensor(1);
+  if (reduction_indices_t == nullptr || !c->RankKnown(input)) {
+    // If we do not have the reduction values at runtime, or the
+    // rank of the input, we don't know the output shape.
+
+    if (keep_dims && c->RankKnown(input)) {
+      // output rank matches input input if <keep_dims>.
+      c->set_output(0, c->UnknownShapeOfRank(c->Rank(input)));
+      return Status::OK();
+    } else {
+      return shape_inference::UnknownShape(c);
+    }
+  }
+
+  const int32 input_rank = c->Rank(input);
+  std::set<int64> true_indices;
+  if (reduction_indices_t->dtype() == DataType::DT_INT32) {
+    TF_RETURN_IF_ERROR(ReductionShapeHelper<int32>(reduction_indices_t,
+                                                   input_rank, true_indices));
+  } else if (reduction_indices_t->dtype() == DataType::DT_INT64) {
+    TF_RETURN_IF_ERROR(ReductionShapeHelper<int64>(reduction_indices_t,
+                                                   input_rank, true_indices));
+  } else {
+    return errors::InvalidArgument(
+        "reduction_indices can only be int32 or int64");
+  }
+
+  std::vector<DimensionHandle> dims;
+  for (int i = 0; i < input_rank; ++i) {
+    if (true_indices.count(i) > 0) {
+      if (keep_dims) {
+        dims.emplace_back(c->MakeDim(1));
+      }
+    } else {
+      dims.emplace_back(c->Dim(input, i));
+    }
+  }
+
+  c->set_output(0, c->MakeShape(dims));
+  return Status::OK();
+}
+
+Status ConcatShapeHelper(InferenceContext* c, int start_value_index,
+                         int end_value_index, int dim_index) {
+  ShapeHandle unused;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(dim_index), 0, &unused));
+  const Tensor* concat_dim_t = c->input_tensor(dim_index);
+  if (concat_dim_t == nullptr) {
+    // Return an unknown shape with same rank as inputs, or an unknown rank
+    // if no input's rank is known.
+
+    // Find rank.
+    int32 rank = InferenceContext::kUnknownRank;
+    for (int i = start_value_index; i < end_value_index; ++i) {
+      if (rank == InferenceContext::kUnknownRank) rank = c->Rank(c->input(i));
+      if (rank != InferenceContext::kUnknownRank) {
+        break;
+      }
+    }
+    if (rank == InferenceContext::kUnknownRank) {
+      c->set_output(0, c->UnknownShape());
+      return Status::OK();
+    } else if (rank == 0) {
+      return errors::InvalidArgument(
+          "Can't concatenate scalars (use tf.stack instead)");
+    } else {
+      for (int i = start_value_index; i < end_value_index; ++i) {
+        // Check that all the inputs are of the correct rank.
+        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), rank, &unused));
+      }
+    }
+    // Build result of <rank> different unknown dims.
+    std::vector<DimensionHandle> dims;
+    dims.reserve(rank);
+    for (int i = 0; i < rank; ++i) dims.push_back(c->UnknownDim());
+    c->set_output(0, c->MakeShape(dims));
+    return Status::OK();
+  }
+
+  // Merge all the non-concat dims, and sum the concat dim to make an output
+  // shape.
+  const int32 concat_dim = concat_dim_t->scalar<int32>()();
+
+  // Minimum required number of dimensions.
+  const int min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
+
+  ShapeHandle output_before;
+  ShapeHandle output_after;
+
+  ShapeHandle input = c->input(end_value_index - 1);
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
+  TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &output_before));
+  DimensionHandle output_middle = c->Dim(input, concat_dim);
+  if (concat_dim == -1) {
+    output_after = c->Scalar();  // no dimensions.
+  } else {
+    TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &output_after));
+  }
+
+  for (int i = end_value_index - 2; i >= start_value_index; --i) {
+    ShapeHandle before;
+    ShapeHandle after;
+    input = c->input(i);
+    TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
+    TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &before));
+    DimensionHandle middle = c->Dim(input, concat_dim);
+    if (concat_dim == -1) {
+      after = c->Scalar();
+    } else {
+      TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &after));
+    }
+
+    TF_RETURN_IF_ERROR(c->Merge(before, output_before, &output_before));
+    TF_RETURN_IF_ERROR(c->Add(output_middle, middle, &output_middle));
+    TF_RETURN_IF_ERROR(c->Merge(after, output_after, &output_after));
+  }
+
+  ShapeHandle s;
+  TF_RETURN_IF_ERROR(
+      c->Concatenate(output_before, c->Vector(output_middle), &s));
+  TF_RETURN_IF_ERROR(c->Concatenate(s, output_after, &s));
+  c->set_output(0, s);
+  return Status::OK();
+}
+
+Status ConcatShape(InferenceContext* c, int num_inputs_to_concat) {
+  return ConcatShapeHelper(c, 1 /* start_value_index */,
+                           1 + num_inputs_to_concat /* end_value_index */,
+                           0 /* dim_index */);
+}
+
+Status ConcatV2Shape(InferenceContext* c) {
+  return ConcatShapeHelper(c, 0 /* start_value_index */,
+                           c->num_inputs() - 1 /* end_value_index */,
+                           c->num_inputs() - 1 /* dim_index */);
+}
+
+Status BroadcastBinaryOpShapeFn(InferenceContext* c) {
+  ShapeHandle shape_x = c->input(0);
+  ShapeHandle shape_y = c->input(1);
+  if (!c->RankKnown(shape_x) || !c->RankKnown(shape_y)) {
+    c->set_output(0, c->UnknownShape());
+    return Status::OK();
+  }
+  const int32 rank_x = c->Rank(shape_x);
+  const int32 rank_y = c->Rank(shape_y);
+  const int32 rank_out = std::max(rank_x, rank_y);
+
+  // To compute the broadcast dimensions, we zip together shape_x and shape_y
+  // and
+  // pad with 1 to make them the same length.
+  std::vector<DimensionHandle> dims;
+  DimensionHandle dim_one;
+  if (rank_x != rank_y) dim_one = c->MakeDim(1);
+  for (int i = 0; i < rank_out; ++i) {
+    const auto dim_x = i < (rank_out - rank_x)
+                           ? dim_one
+                           : c->Dim(shape_x, i - (rank_out - rank_x));
+    const bool dim_y_is_one = (i < (rank_out - rank_y));
+    const auto dim_y =
+        dim_y_is_one ? dim_one : c->Dim(shape_y, i - (rank_out - rank_y));
+    if (!c->ValueKnown(dim_x) || !c->ValueKnown(dim_y)) {
+      // One or both dimensions is unknown.
+      //
+      // - If either dimension is greater than 1, we assume that the program is
+      // correct, and the other dimension will be broadcast to match it.
+      // TODO(cwhipkey): For shape inference, if we eliminate the shape checks
+      // in C++ op code, we must still assert that the unknown dim is either 1
+      // or the same as the known dim.
+      // - If either dimension is 1, the other dimension is the output.
+      if (c->Value(dim_x) > 1) {
+        dims.push_back(dim_x);
+      } else if (c->Value(dim_y) > 1) {
+        dims.push_back(dim_y);
+      } else if (c->Value(dim_x) == 1) {
+        dims.push_back(dim_y);
+      } else if (c->Value(dim_y) == 1) {
+        dims.push_back(dim_x);
+      } else if (dim_y.SameHandle(dim_x)) {
+        dims.push_back(dim_x);
+      } else {
+        dims.push_back(c->UnknownDim());
+      }
+    } else if (c->Value(dim_x) == 1 || c->Value(dim_y) == 1) {
+      if (c->Value(dim_x) == 1 && !dim_y_is_one) {
+        // We will broadcast dim_x to dim_y.
+        dims.push_back(dim_y);
+      } else {
+        DCHECK_EQ(c->Value(dim_y), 1);
+        // We will broadcast dim_y to dim_x.
+        dims.push_back(dim_x);
+      }
+    } else {
+      DimensionHandle dim;
+      TF_RETURN_IF_ERROR(c->Merge(dim_x, dim_y, &dim));
+      dims.push_back(dim);
+    }
+  }
+
+  c->set_output(0, c->MakeShape(dims));
+  return Status::OK();
+}
+
+Status RandomShape(shape_inference::InferenceContext* c) {
+  shape_inference::ShapeHandle out;
+  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out));
+  c->set_output(0, out);
+  return Status::OK();
+}
+
+Status ValidateSparseTensor(InferenceContext* c, ShapeHandle indices_shape,
+                            ShapeHandle values_shape, ShapeHandle shape_shape) {
+  // Validate ranks.
+  ShapeHandle unused_shape;
+  TF_RETURN_IF_ERROR(c->WithRank(indices_shape, 2, &unused_shape));
+  TF_RETURN_IF_ERROR(c->WithRank(values_shape, 1, &unused_shape));
+  TF_RETURN_IF_ERROR(c->WithRank(shape_shape, 1, &unused_shape));
+
+  // Number of elements in indices and values must match.
+  DimensionHandle num_index_elements_dim = c->Dim(indices_shape, 0);
+  if (c->ValueKnown(num_index_elements_dim)) {
+    DimensionHandle num_values_elements_dim = c->Dim(values_shape, 0);
+    if (c->ValueKnown(num_values_elements_dim)) {
+      int64 num_index_elements = c->Value(num_index_elements_dim);
+      int64 num_values_elements = c->Value(num_values_elements_dim);
+      if (num_index_elements != num_values_elements) {
+        return errors::InvalidArgument("Number of elements in index (",
+                                       num_index_elements, ") and values (",
+                                       num_values_elements, ") do not match.");
+      }
+    }
+  }
+
+  // Rank embedded in indices must match shape.
+  DimensionHandle index_rank_dim = c->Dim(indices_shape, 1);
+  if (c->ValueKnown(index_rank_dim)) {
+    DimensionHandle shape_rank_dim = c->Dim(shape_shape, 0);
+    if (c->ValueKnown(shape_rank_dim)) {
+      int64 index_rank = c->Value(index_rank_dim);
+      int32 shape_rank = c->Value(shape_rank_dim);
+      if (index_rank != shape_rank) {
+        return errors::InvalidArgument("Index rank (", index_rank,
+                                       ") and shape rank (", shape_rank,
+                                       ") do not match.");
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+Status ScatterNdUpdateShape(InferenceContext* c) {
+  ShapeHandle input_shape = c->input(0);
+  if (c->input_handle_shapes_and_types(0) != nullptr) {
+    input_shape = (*c->input_handle_shapes_and_types(0))[0].shape;
+  }
+  ShapeHandle indices_shape;
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &indices_shape));
+  ShapeHandle updates_shape;
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(2), 1, &updates_shape));
+
+  if (c->Value(c->NumElements(input_shape)) == 0 &&
+      (c->Value(c->NumElements(indices_shape)) > 0 ||
+       c->Value(c->NumElements(updates_shape)) > 0)) {
+    return errors::InvalidArgument(
+        "Indices and updates specified for empty output shape");
+  }
+
+  if (c->RankKnown(indices_shape) && c->RankKnown(updates_shape)) {
+    const int64 num_outer_dims = c->Rank(indices_shape) - 1;
+    const DimensionHandle index_size = c->Dim(indices_shape, -1);
+
+    // We can only do more validation if the last dimension of indices
+    // is a known value.
+    if (c->ValueKnown(index_size)) {
+      const int64 ix = c->Value(index_size);
+      ShapeHandle unused;
+      ShapeHandle prefix_indices;
+      TF_RETURN_IF_ERROR(
+          c->Subshape(indices_shape, 0, num_outer_dims, &prefix_indices));
+      ShapeHandle prefix_updates;
+      TF_RETURN_IF_ERROR(
+          c->Subshape(updates_shape, 0, num_outer_dims, &prefix_updates));
+
+      Status s = c->Merge(prefix_indices, prefix_updates, &unused);
+      if (!s.ok()) {
+        return errors::InvalidArgument(
+            "The outer ", num_outer_dims, " dimensions of indices.shape=",
+            c->DebugString(indices_shape), " must match the outer ",
+            num_outer_dims, " dimensions of updates.shape=",
+            c->DebugString(updates_shape), ": ", s.error_message());
+      }
+
+      ShapeHandle input_suffix;
+      TF_RETURN_IF_ERROR(c->Subshape(input_shape, ix, &input_suffix));
+      ShapeHandle suffix_updates;
+      TF_RETURN_IF_ERROR(
+          c->Subshape(updates_shape, num_outer_dims, &suffix_updates));
+      s = c->Merge(input_suffix, suffix_updates, &unused);
+      if (!s.ok()) {
+        return errors::InvalidArgument(
+            "The inner ", c->Rank(input_shape) - ix,
+            " dimensions of input.shape=", c->DebugString(input_shape),
+            " must match the inner ", c->Rank(updates_shape) - num_outer_dims,
+            " dimensions of updates.shape=", c->DebugString(updates_shape),
+            ": ", s.error_message());
+      }
+    }
+  }
+
+  if (c->input_handle_shapes_and_types(0) == nullptr) {
+    c->set_output(0, input_shape);
+  }
+  return Status::OK();
+}
+
+Status ExplicitShape(InferenceContext* c) {
+  PartialTensorShape shape;
+  TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape));
+  ShapeHandle output_shape;
+  TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape, &output_shape));
+  c->set_output(0, output_shape);
+  return Status::OK();
+}
+
+}  // namespace shape_inference
+
+}  // namespace tensorflow

common_shape_fns.h

@@ -0,0 +1,290 @@
+/* Copyright 2016 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.
+==============================================================================*/
+#ifndef THIRD_PARTY_TENSORFLOW_CORE_OPS_COMMON_SHAPE_FNS_H_
+#define THIRD_PARTY_TENSORFLOW_CORE_OPS_COMMON_SHAPE_FNS_H_
+
+#include <array>
+
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+// GetWindowedOutputSize(): Given an input tensor, kernel, stride and padding
+// type, the function computes the output and padding dimensions.
+//
+// For example, ignoring batches or multiple features, a 1D convolution
+// takes as input a 1D tensor of shape (H), and convolves it with a filter of
+// shape (K).
+//
+// It also takes in a few additional parameters:
+//
+// Stride (S): the stride with which we apply the filters. This is the offset
+// between locations where we apply the filters. A larger stride
+// means that the output will be spatially smaller.
+//
+// Padding (P): the padding we apply to the input tensor along each
+// dimension. This is usually used to make sure that the spatial dimensions
+// do not shrink when we progress with convolutions. Two types of padding are
+// often used:
+//   SAME: the pad value is computed so that the output will have size H/S.
+//   VALID: no padding is carried out.
+// The padded area is zero-filled.
+//
+// The output dimensions for convolution and many other operations, when given
+// all the parameters above, are as follows:
+// - When Padding = SAME: the output size is (H'), where
+//     H' = ceil(float(H) / float(S))
+//   where ceil is the ceiling function. The number of padded cells
+//   is computed as:
+//     Pc = ((H' - 1) * S + K - H) / 2
+//   When the stride is 1, the expression simplifies to
+//     H' = H, Pc = (K-1)/2.
+//   This is where SAME comes from - the output has the same size as the input
+//   has.
+//
+// - When Padding = VALID: the output size is computed as
+//     H' = ceil(float(H - K + 1) / float(S))
+//   and the number of padded cells is always zero.
+//   When the stride is 1, the expression simplifies to
+//     H' = H-K+1.
+//
+// For convolution, mathematically, the output value at location (r')
+// is the inner product of two vectors: the chunk of input at
+//    ((r'*S-Pr) : (r'*S-Pr+K)),
+// and the filter.
+//
+// For 2D and 3D convolutions, the spatial dimensions are orthogonal, so the
+// size and padding of each spatial dimension can be computed by calling
+// GetWindowedOutputSize separately for each dimension.
+//
+Status GetWindowedOutputSize(int64 input_size, int64 filter_size, int64 stride,
+                             Padding padding_type, int64* output_size,
+                             int64* padding_size);
+
+// The V2 version computes the same outputs with arbitrary dilation_rate.
+// The output dimensions are computed as follows:
+// - When adding dilation_rate (D), we compute an effective filter size (K'):
+//     K' = (K - 1) * D + 1
+// - When Padding = SAME: the output size is (H'), where
+//     H' = ceil(float(H) / float(S))
+//   where ceil is the ceiling function. The number of padded cells
+//   is computed as:
+//     Pc = ((H' - 1) * S + K' - H) / 2
+//   When the stride is 1, the expression simplifies to
+//     H' = H, Pc = (K'-1)/2.
+//   This is where SAME comes from - the output has the same size as the input
+//   has.
+//
+// - When Padding = VALID: the output size is computed as
+//     H' = ceil(float(H - K' + 1) / float(S))
+//   and the number of padded cells is always zero.
+//   When the stride is 1, the expression simplifies to
+//     H' = H-K'+1.
+//
+// TODO(b/67112639): Merge V2 versions and the original versions eventually.
+Status GetWindowedOutputSizeV2(int64 input_size, int64 filter_size,
+                               int64 dilation_rate, int64 stride,
+                               Padding padding_type, int64* output_size,
+                               int64* padding_size);
+
+// Returns the same output dimensions as in GetWindowedOutputSize, but returns
+// verbose padding dimensions (before/after). Any excess padding
+// (caused by an odd padding size value) is added to the 'padding_after'
+// dimension.
+Status GetWindowedOutputSizeVerbose(int64 input_size, int64 filter_size,
+                                    int64 stride, Padding padding_type,
+                                    int64* output_size, int64* padding_before,
+                                    int64* padding_after);
+
+// The V2 version computes the same outputs with arbitrary dilation_rate. For
+// detailed equations, refer to the comments for GetWindowedOutputSizeV2().
+Status GetWindowedOutputSizeVerboseV2(int64 input_size, int64 filter_size,
+                                      int64 dilation_rate, int64 stride,
+                                      Padding padding_type, int64* output_size,
+                                      int64* padding_before,
+                                      int64* padding_after);
+
+// Given an input tensor, kernel, stride and padding type, populates the 3D size
+// of the output tensor and padding to be applied to the input tensor at the
+// lower end of every dimension. Use for 3D convolutions, where the input data
+// is padded with zeros, as well as for 3D avg/max pooling, where the input data
+// is padded with invalid values that are not considered for pooling.
+Status Get3dOutputSize(const std::array<int64, 3>& input,
+                       const std::array<int64, 3>& window,
+                       const std::array<int64, 3>& strides,
+                       Padding padding_type, std::array<int64, 3>* output_ptr,
+                       std::array<int64, 3>* padding_ptr);
+
+// The V2 version computes the same outputs with arbitrary dilation_rate. For
+// detailed equations, refer to the comments for GetWindowedOutputSizeV2().
+Status Get3dOutputSizeV2(const std::array<int64, 3>& input,
+                         const std::array<int64, 3>& window,
+                         const std::array<int64, 3>& dilations,
+                         const std::array<int64, 3>& strides,
+                         Padding padding_type, std::array<int64, 3>* output_ptr,
+                         std::array<int64, 3>* padding_ptr);
+
+namespace shape_inference {
+
+// Like GetWindowedOutputSize, but deals with DimensionHandles.
+Status GetWindowedOutputSizeFromDims(InferenceContext* c,
+                                     DimensionHandle input_size,
+                                     DimensionOrConstant filter_size,
+                                     int64 stride, Padding padding_type,
+                                     DimensionHandle* output_size);
+
+// The V2 version computes the same outputs with arbitrary dilation_rate. For
+// detailed equations, refer to the comments for GetWindowedOutputSizeV2().
+Status GetWindowedOutputSizeFromDimsV2(InferenceContext* c,
+                                       DimensionHandle input_size,
+                                       DimensionOrConstant filter_size,
+                                       int64 dilation_rate, int64 stride,
+                                       Padding padding_type,
+                                       DimensionHandle* output_size);
+
+// Transfers shape of input(0) to output(0).
+Status UnchangedShape(shape_inference::InferenceContext* c);
+
+// Transfers shape of input(0) to output(0), after asserting its rank is <rank>.
+inline Status UnchangedShapeWithRank(shape_inference::InferenceContext* c,
+                                     int32 rank) {
+  ShapeHandle out;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &out));
+  c->set_output(0, out);
+  return Status::OK();
+}
+
+// Transfers shape of input(0) to output(0), after asserting its rank >= <rank>.
+inline Status UnchangedShapeWithRankAtLeast(
+    shape_inference::InferenceContext* c, int32 rank) {
+  ShapeHandle out;
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), rank, &out));
+  c->set_output(0, out);
+  return Status::OK();
+}
+
+// Transfers shape of input(0) to output(0), after asserting its rank <= <rank>.
+inline Status UnchangedShapeWithRankAtMost(shape_inference::InferenceContext* c,
+                                           int32 rank) {
+  ShapeHandle out;
+  TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), rank, &out));
+  c->set_output(0, out);
+  return Status::OK();
+}
+
+// Shape function for use with ops no outputs.
+inline Status NoOutputs(shape_inference::InferenceContext* c) {
+  return Status::OK();
+}
+
+// Shape function for ops that output a single scalar value.
+inline Status ScalarShape(shape_inference::InferenceContext* c) {
+  c->set_output(0, c->Scalar());
+  return Status::OK();
+}
+
+// Shape function for binary ops where both inputs and the output match.
+inline Status MergeBothInputsShapeFn(InferenceContext* c) {
+  ShapeHandle out;
+  TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(1), &out));
+  c->set_output(0, out);
+  return Status::OK();
+}
+
+// Returns a new shape with the specified dims arranged in the specified
+// format. The returned value is owned by this context.
+// Note: if format = "FORMAT_NCHW_VECT_C" then C represents the outer_depth.
+Status MakeShapeFromFormat(TensorFormat format, DimensionOrConstant N,
+                           const std::vector<DimensionOrConstant>& spatial,
+                           DimensionOrConstant C, ShapeHandle* out,
+                           shape_inference::InferenceContext* context);
+
+// Shape function for MatMul-like operations.
+Status MatMulShape(shape_inference::InferenceContext* c);
+
+// Shape function for BiasAdd-like operations.
+Status BiasAddShape(shape_inference::InferenceContext* c);
+
+// Shape function for BiasAddGrad-like operations.
+Status BiasAddGradShape(shape_inference::InferenceContext* c);
+
+// Shape function for Conv2D-like operations.
+Status Conv2DShape(shape_inference::InferenceContext* c);
+
+// Shape function for Conv3D-like operations.
+Status Conv3DShape(shape_inference::InferenceContext* c);
+
+// Shape function for DepthwiseConv2D-like operations.
+Status DepthwiseConv2DNativeShape(shape_inference::InferenceContext* c);
+
+// Shape function for AvgPool-like operations.
+Status AvgPoolShape(shape_inference::InferenceContext* c);
+
+// Shape function for FusedBatchNorm and FusedBatchNormV2 operations.
+Status FusedBatchNormShape(shape_inference::InferenceContext* c);
+
+// Shape function for FusedBatchNormGrad and FusedBatchNormGradV2 operations.
+Status FusedBatchNormGradShape(shape_inference::InferenceContext* c);
+
+// Shape function for MaxPool-like operations.
+Status MaxPoolShape(shape_inference::InferenceContext* c);
+
+// Shape function for MaxPoolV2-like operations.
+Status MaxPoolV2Shape(shape_inference::InferenceContext* c, int num_inputs);
+
+// Shape function for 3D Pooling operations.
+Status Pool3DShape(shape_inference::InferenceContext* c);
+
+// Shape function for use with ops whose output shapes are unknown.
+Status UnknownShape(shape_inference::InferenceContext* c);
+
+// Shape function for reduction operations.
+Status ReductionShape(shape_inference::InferenceContext* c);
+
+// Shape function for concat operations.
+// <num_inputs_to_concat> is the number of inputs to concatenate and are taken
+// from inputs
+// [1,num_inputs_to_concat] of the op.  Input 0 is the concat_dim input.
+Status ConcatShape(shape_inference::InferenceContext* c,
+                   int num_inputs_to_concat);
+
+// Shape function for concat operations.
+Status ConcatV2Shape(shape_inference::InferenceContext* c);
+
+// Shape function for binary operators that broadcast their inputs.
+// Tested by ops/math_ops_test.cc.
+Status BroadcastBinaryOpShapeFn(InferenceContext* c);
+
+// Shape function for random operations.
+Status RandomShape(shape_inference::InferenceContext* c);
+
+// Validates the 3 component tensors of a sparse tensor have the proper
+// shapes. This mimics SparseTensor.__init__ in python/framework/ops.py.
+Status ValidateSparseTensor(InferenceContext* c, ShapeHandle indices_shape,
+                            ShapeHandle values_shape, ShapeHandle shape_shape);
+
+// Shape function for ScatterNd update/add/sub/... operations.
+Status ScatterNdUpdateShape(InferenceContext* c);
+
+// Shape function for ops with an explicit "shape" attribute.
+Status ExplicitShape(InferenceContext* c);
+
+}  // namespace shape_inference
+
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_TENSORFLOW_CORE_OPS_COMMON_SHAPE_FNS_H_

common_shape_fns_test.cc

@@ -0,0 +1,1131 @@
+/* Copyright 2016 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.
+==============================================================================*/
+#include "tensorflow/core/framework/common_shape_fns.h"
+
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/shape_inference_testutil.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace shape_inference {
+
+namespace {
+
+PartialTensorShape S(std::initializer_list<int64> dims) {
+  return PartialTensorShape(dims);
+}
+
+PartialTensorShape Unknown() { return PartialTensorShape(); }
+
+OpDef MakeOpDef(int num_inputs, int num_outputs) {
+  OpRegistrationData op_reg_data;
+  OpDefBuilder b("dummy");
+  for (int i = 0; i < num_inputs; ++i) {
+    b.Input(strings::StrCat("i", i, ": float"));
+  }
+  for (int i = 0; i < num_outputs; ++i) {
+    b.Output(strings::StrCat("o", i, ": float"));
+  }
+  CHECK(b.Attr("foo:string").Finalize(&op_reg_data).ok());
+  return op_reg_data.op_def;
+}
+
+}  // namespace
+
+TEST(CommonShapeFnsTest, NoOutputShapeTest) {
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(OpDefBuilder("Assert")
+                  .Input("condition: bool")
+                  .Input("data: float")
+                  .Finalize(&op_reg_data));
+  OpDef op_def = op_reg_data.op_def;
+
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("test", "Assert")
+                  .Input("condition", 0, DT_BOOL)
+                  .Input({{"data", 0, DT_FLOAT}})
+                  .Finalize(&def));
+
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({}), S({10})}, {},
+                     {}, {});
+  TF_EXPECT_OK(NoOutputs(&c));
+  EXPECT_EQ(0, c.num_outputs());
+}
+
+TEST(CommonShapeFnsTest, ScalarShapeTest) {
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(OpDefBuilder("L2Loss")
+                  .Input("t: float")
+                  .Output("t: float")
+                  .Finalize(&op_reg_data));
+  OpDef op_def = op_reg_data.op_def;
+
+  NodeDef def;
+  TF_CHECK_OK(
+      NodeDefBuilder("test", "L2Loss").Input("t", 0, DT_FLOAT).Finalize(&def));
+
+  {
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({})}, {}, {}, {});
+    TF_EXPECT_OK(ScalarShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(0, c.Rank(output));
+  }
+
+  {
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({1, 23, 4, 4, 2})}, {}, {}, {});
+    TF_EXPECT_OK(ScalarShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(0, c.Rank(output));
+  }
+}
+
+TEST(CommonShapeFnsTest, MatMulShapeTest) {
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(OpDefBuilder("MatMul")
+                  .Input("a: float")
+                  .Input("b: float")
+                  .Output("c: float")
+                  .Attr("transpose_a:bool=false")
+                  .Attr("transpose_b:bool=false")
+                  .Finalize(&op_reg_data));
+  OpDef op_def = op_reg_data.op_def;
+
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("test", "MatMul")
+                  .Input("a", 0, DT_FLOAT)
+                  .Input("b", 0, DT_FLOAT)
+                  .Attr("transpose_a", false)
+                  .Attr("transpose_b", false)
+                  .Finalize(&def));
+
+  {
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 3}), S({3, 4})}, {}, {}, {});
+    TF_EXPECT_OK(MatMulShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_EQ(4, c.Value(c.Dim(output, 1)));
+  }
+
+  {
+    // Unknown inner dimension for one
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, -1}), S({3, 4})}, {}, {}, {});
+    TF_EXPECT_OK(MatMulShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_EQ(4, c.Value(c.Dim(output, 1)));
+  }
+
+  {
+    // Invalid rank.
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({2}), S({3, 4})},
+                       {}, {}, {});
+    auto s = MatMulShape(&c);
+    EXPECT_FALSE(s.ok());
+    EXPECT_TRUE(
+        StringPiece(s.ToString())
+            .contains("Invalid argument: Shape must be rank 2 but is rank 1"));
+  }
+
+  {
+    // Unknown outer dimension
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 3}), S({3, -1})}, {}, {}, {});
+    TF_EXPECT_OK(MatMulShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_FALSE(c.ValueKnown(c.Dim(output, 1)));
+  }
+
+  {
+    // Inner shapes not compatible
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 5}), S({3, 4})}, {}, {}, {});
+    auto s = MatMulShape(&c);
+    EXPECT_FALSE(s.ok());
+    EXPECT_TRUE(
+        StringPiece(s.ToString())
+            .contains(
+                "Invalid argument: Dimensions must be equal, but are 5 and 3"));
+  }
+
+  {
+    // Inner shapes not compatible
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 5, 3}), S({3, 5, 4})}, {}, {}, {});
+    auto s = MatMulShape(&c);
+    EXPECT_FALSE(s.ok());
+    EXPECT_TRUE(
+        StringPiece(s.ToString())
+            .contains("Invalid argument: Shape must be rank 2 but is rank 3"));
+  }
+
+  {
+    // transpose_a
+    TF_CHECK_OK(NodeDefBuilder("test", "MatMul")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("transpose_a", true)
+                    .Attr("transpose_b", false)
+                    .Attr("type", DT_FLOAT)
+                    .Finalize(&def));
+
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({3, 2}), S({3, 4})}, {}, {}, {});
+    auto s = MatMulShape(&c);
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_EQ(4, c.Value(c.Dim(output, 1)));
+  }
+
+  {
+    // transpose_b
+    TF_CHECK_OK(NodeDefBuilder("test", "MatMul")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("transpose_a", false)
+                    .Attr("transpose_b", true)
+                    .Attr("type", DT_FLOAT)
+                    .Finalize(&def));
+
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 3}), S({4, 3})}, {}, {}, {});
+    auto s = MatMulShape(&c);
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_EQ(4, c.Value(c.Dim(output, 1)));
+  }
+}
+
+TEST(CommonShapeFnsTest, BiasAddShapeTest) {
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(OpDefBuilder("BiasAdd")
+                  .Input("a: float")
+                  .Input("b: float")
+                  .Output("c: float")
+                  .Finalize(&op_reg_data));
+
+  OpDef op_def = op_reg_data.op_def;
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("test", "BiasAdd")
+                  .Input("a", 0, DT_FLOAT)
+                  .Input("b", 0, DT_FLOAT)
+                  .Finalize(&def));
+
+  {
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 10}), S({10})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(2, c.Value(c.Dim(output, 0)));
+    EXPECT_EQ(10, c.Value(c.Dim(output, 1)));
+  }
+
+  {
+    // Unknown ranks.
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {Unknown(), Unknown()}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_FALSE(c.RankKnown(output));
+  }
+
+  {
+    // Rank > 2
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({4, 3, 4, 2, 15}), S({15})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ("[4,3,4,2,15]", c.DebugString(output));
+  }
+
+  {
+    // NCHW format
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAdd")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({2, 3, 4, 5}), S({3})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ("[2,3,4,5]", c.DebugString(output));
+  }
+
+  {
+    // NCHW format with high input rank
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAdd")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({8, 6, 4, 2, 3, 4, 5}), S({3})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ("[8,6,4,2,3,4,5]", c.DebugString(output));
+  }
+
+  {
+    // NCHW format with input rank 3
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAdd")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({10, 11, 12}), S({10})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ("[10,11,12]", c.DebugString(output));
+  }
+
+  {
+    // Input rank not high enough
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({3}), S({3})}, {},
+                       {}, {});
+    EXPECT_FALSE(BiasAddShape(&c).ok());
+  }
+
+  {
+    // NCHW rank not high enough
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAdd")
+                    .Input("a", 0, DT_FLOAT)
+                    .Input("b", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    // NCHW format
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({2, 3}), S({3})},
+                       {}, {}, {});
+    EXPECT_FALSE(BiasAddShape(&c).ok());
+  }
+}
+
+TEST(CommonShapeFnsTest, BiasAddGradShapeTest) {
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(OpDefBuilder("BiasAddGrad")
+                  .Input("a: float")
+                  .Output("b: float")
+                  .Finalize(&op_reg_data));
+
+  OpDef op_def = op_reg_data.op_def;
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("test", "BiasAddGrad")
+                  .Input("a", 0, DT_FLOAT)
+                  .Finalize(&def));
+
+  {
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({2, 10})}, {}, {},
+                       {});
+    TF_EXPECT_OK(BiasAddGradShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(10, c.Value(c.Dim(output, 0)));
+  }
+
+  {
+    // Rank > 2
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({5, 7, 2, 10})},
+                       {}, {}, {});
+    TF_EXPECT_OK(BiasAddGradShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(10, c.Value(c.Dim(output, 0)));
+  }
+
+  {
+    // NCHW format
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAddGrad")
+                    .Input("a", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({2, 3, 4, 5})},
+                       {}, {}, {});
+    TF_EXPECT_OK(BiasAddGradShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(3, c.Value(c.Dim(output, 0)));
+  }
+
+  {
+    // NCHW format with high input rank
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAddGrad")
+                    .Input("a", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def,
+                       {S({8, 6, 4, 2, 3, 4, 5})}, {}, {}, {});
+    TF_EXPECT_OK(BiasAddGradShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(3, c.Value(c.Dim(output, 0)));
+  }
+
+  {
+    // NCHW format with input rank 3
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAddGrad")
+                    .Input("a", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({10, 11, 12})},
+                       {}, {}, {});
+    TF_EXPECT_OK(BiasAddGradShape(&c));
+    ShapeHandle output = c.output(0);
+    EXPECT_EQ(10, c.Value(c.Dim(output, 0)));
+  }
+
+  {
+    // Input rank not high enough
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({3})}, {}, {},
+                       {});
+    EXPECT_FALSE(BiasAddGradShape(&c).ok());
+  }
+
+  {
+    // NCHW rank not high enough
+    TF_CHECK_OK(NodeDefBuilder("test", "BiasAddGrad")
+                    .Input("a", 0, DT_FLOAT)
+                    .Attr("data_format", "NCHW")
+                    .Finalize(&def));
+    // NCHW format
+    InferenceContext c(TF_GRAPH_DEF_VERSION, &def, op_def, {S({2, 3})}, {}, {},
+                       {});
+    EXPECT_FALSE(BiasAddGradShape(&c).ok());
+  }
+}
+
+TEST(CommonShapeFnsTest, Conv2DShapeTest) {
+  ShapeInferenceTestOp op("Conv2D");
+  auto set_op = [&op](const std::vector<int32>& strides, const string& padding,
+                      const string& data_format, const string& filter_format) {
+    TF_CHECK_OK(NodeDefBuilder("test", "Conv2D")
+                    .Input("input", 0, DT_FLOAT)
+                    .Input("filter", 0, DT_FLOAT)
+                    .Attr("strides", strides)
+                    .Attr("padding", padding)
+                    .Attr("data_format", data_format)
+                    .Attr("filter_format", filter_format)
+                    .Finalize(&op.node_def));
+  };
+
+  // Invalid rank for input
+  INFER_ERROR("must be rank 4", op, "[4,4];[2,1,1,1]");
+  // Invalid rank for filter
+  INFER_ERROR("must be rank 4", op, "[1,4,4,1];[2,1,1]");
+
+  // Invalid value for strides
+  set_op({{1, 1, 0, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_ERROR("must be > 0", op, "[1,2,2,1];[1,1,1,1]");
+
+  // 1x1 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_OK(op, "[1,2,2,1];[1,1,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 2x2 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_OK(op, "[1,2,2,1];[2,2,1,1]", "[d0_0,1,1,d1_3]");
+
+  // 3x3 input, 1x1 filter, 2x2 stride
+  set_op({{1, 2, 2, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_OK(op, "[1,3,3,1];[1,1,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 3x3 input, 1x1 filter, 2x1 stride
+  set_op({{1, 2, 1, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_OK(op, "[1,3,3,1];[1,1,1,1]", "[d0_0,2,3,d1_3]");
+
+  // 4x4 input, 2x1 filter, 1x2 stride
+  set_op({{1, 1, 2, 1}}, "VALID", "NHWC", "HWIO");
+  INFER_OK(op, "[1,4,4,1];[2,1,1,1]", "[d0_0,3,2,d1_3]");
+
+  // Unknown dims in the critical fields lead to partial inference.
+  INFER_OK(op, "[1,4,4,1];[2,1,1,1]", "[d0_0,3,2,d1_3]");
+  INFER_OK(op, "[1,?,4,1];[2,1,1,1]", "[d0_0,?,2,d1_3]");
+  INFER_OK(op, "[1,4,?,1];[2,1,1,1]", "[d0_0,3,?,d1_3]");
+  INFER_OK(op, "[1,4,4,?];[2,1,1,1]", "[d0_0,3,2,d1_3]");
+  INFER_OK(op, "[1,4,4,1];[?,1,1,1]", "[d0_0,?,2,d1_3]");
+  INFER_OK(op, "[1,4,4,1];[2,?,1,1]", "[d0_0,3,?,d1_3]");
+
+  // input depths must match.
+  INFER_ERROR("Dimensions must be equal, but are 10 and 10000", op,
+              "[1,2,2,10];[1,1,10000,20]");
+
+  // Tests for NCHW
+  // 1x1 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NCHW", "HWIO");
+  INFER_OK(op, "[1,1,2,2];[1,1,1,1]", "[d0_0,d1_3,2,2]");
+
+  // 2x2 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NCHW", "HWIO");
+  INFER_OK(op, "[1,1,2,2];[2,2,1,1]", "[d0_0,d1_3,1,1]");
+
+  // 3x3 input, 1x1 filter, 2x2 stride
+  set_op({{1, 1, 2, 2}}, "VALID", "NCHW", "HWIO");
+  INFER_OK(op, "[1,1,3,3];[1,1,1,1]", "[d0_0,d1_3,2,2]");
+
+  // 3x3 input, 1x1 filter, 2x1 stride
+  set_op({{1, 1, 2, 1}}, "VALID", "NCHW", "HWIO");
+  INFER_OK(op, "[1,1,3,3];[1,1,1,1]", "[d0_0,d1_3,2,3]");
+
+  // 4x4 input, 2x1 filter, 1x2 stride
+  set_op({{1, 1, 1, 2}}, "VALID", "NCHW", "HWIO");
+  INFER_OK(op, "[1,1,4,4];[2,1,1,1]", "[d0_0,d1_3,3,2]");
+
+  // Tests for NCHW_VECT_C
+  // 1x1 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NCHW_VECT_C", "OIHW_VECT_I");
+  INFER_OK(op, "[1,1,2,2,4];[4,1,1,1,4]", "[d0_0,1,2,2,4]");
+
+  // 2x2 filter
+  set_op({{1, 1, 1, 1}}, "VALID", "NCHW_VECT_C", "OIHW_VECT_I");
+  INFER_OK(op, "[1,1,2,2,4];[4,1,2,2,4]", "[d0_0,1,1,1,4]");
+
+  // 3x3 input, 1x1 filter, 2x2 stride
+  set_op({{1, 1, 2, 2}}, "VALID", "NCHW_VECT_C", "OIHW_VECT_I");
+  INFER_OK(op, "[1,1,3,3,4];[8,1,1,1,4]", "[d0_0,2,2,2,4]");
+
+  // 3x3 input, 1x1 filter, 2x1 stride
+  set_op({{1, 1, 2, 1}}, "VALID", "NCHW_VECT_C", "OIHW_VECT_I");
+  INFER_OK(op, "[1,1,3,3,4];[4,1,1,1,4]", "[d0_0,1,2,3,4]");
+
+  // 4x4 input, 2x1 filter, 1x2 stride
+  set_op({{1, 1, 1, 2}}, "VALID", "NCHW_VECT_C", "OIHW_VECT_I");
+  INFER_OK(op, "[1,1,4,4,4];[4,1,2,1,4]", "[d0_0,1,3,2,4]");
+
+  // Some tests for "SAME" padding
+
+  // 4x4 input, 1x1 filter, 1x1 stride
+  set_op({{1, 1, 1, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[1,4,4,1];[1,1,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // 3x3 input, 2x2 filter, 1x1 stride
+  set_op({{1, 1, 1, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[1,3,3,1];[2,2,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // 4x4 input, 2x2 filter, 2x2 stride
+  set_op({{1, 2, 2, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[1,4,4,1];[2,2,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 4x4 input, 2x2 filter, 1x1 stride
+  set_op({{1, 1, 1, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[1,4,4,1];[2,2,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // With stride 1x1 and SAME, unknown dims don't matter - filter dims except
+  // for output channels are ignored for output, so all inputs are carried
+  // through to output.
+  set_op({{1, 1, 1, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[1,4,4,1];[?,?,?,?]", "[d0_0,d0_1,d0_2,d1_3]");
+  INFER_OK(op, "[1,?,4,1];[?,?,?,?]", "[d0_0,d0_1,d0_2,d1_3]");
+  INFER_OK(op, "[1,4,?,1];[?,?,?,?]", "[d0_0,d0_1,d0_2,d1_3]");
+  INFER_OK(op, "[1,4,4,?];[?,?,?,?]", "[d0_0,d0_1,d0_2,d1_3]");
+  INFER_OK(op, "[?,4,4,1];[?,?,?,?]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // With stride != 1, the input HW dims are divided to produce output dims.
+  set_op({{1, 2, 2, 1}}, "SAME", "NHWC", "HWIO");
+  INFER_OK(op, "[?,4,4,1];[?,?,?,?]", "[d0_0,2,2,d1_3]");
+  INFER_OK(op, "[1,?,4,1];[?,?,?,?]", "[d0_0,?,2,d1_3]");
+  INFER_OK(op, "[1,4,?,1];[?,?,?,?]", "[d0_0,2,?,d1_3]");
+  INFER_OK(op, "[1,4,4,?];[?,?,?,?]", "[d0_0,2,2,d1_3]");
+}
+
+TEST(CommonShapeFnsTest, Conv2DDilatedShapeTest) {
+  ShapeInferenceTestOp op("Conv2D");
+  auto set_op = [&op](const std::vector<int32>& dilations,
+                      const std::vector<int32>& strides, const string& padding,
+                      const string& data_format) {
+    TF_CHECK_OK(NodeDefBuilder("test", "Conv2D")
+                    .Input("input", 0, DT_FLOAT)
+                    .Input("filter", 0, DT_FLOAT)
+                    .Attr("dilations", dilations)
+                    .Attr("strides", strides)
+                    .Attr("padding", padding)
+                    .Attr("data_format", data_format)
+                    .Finalize(&op.node_def));
+  };
+
+  // Invalid rank for dilation
+  set_op({{1, 2, 1}}, {{1, 1, 1, 1}}, "VALID", "NHWC");
+  INFER_ERROR("contain 4 values", op, "[1,2,2,1];[1,1,1,1]");
+
+  // Invalid value for dilation
+  set_op({{1, 0, 1, 1}}, {{1, 1, 1, 1}}, "VALID", "NHWC");
+  INFER_ERROR("must be >= 1", op, "[1,2,2,1];[1,1,1,1]");
+
+  // Tests for NHWC
+  // 1x1 filter, 2x1 dilations, 1x1 strides
+  set_op({{1, 2, 1, 1}}, {{1, 1, 1, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,2,2,1];[1,1,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 1x1 filter, 2x1 dilations, 2x1 strides
+  set_op({{1, 2, 1, 1}}, {{1, 2, 1, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,4,4,1];[1,1,1,1]", "[d0_0,2,4,d1_3]");
+
+  // 1x1 filter, 2x1 dilations, 2x2 strides
+  set_op({{1, 2, 1, 1}}, {{1, 2, 2, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,4,4,1];[1,1,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 3x3 filter, 2x1 dilations, 1x1 strides
+  set_op({{1, 2, 1, 1}}, {{1, 1, 1, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,5,5,1];[3,3,1,1]", "[d0_0,1,3,d1_3]");
+
+  // 3x3 filter, 2x1 dilations, 2x1 strides
+  set_op({{1, 2, 1, 1}}, {{1, 2, 1, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,5,5,1];[3,3,1,1]", "[d0_0,1,3,d1_3]");
+
+  // 3x3 filter, 1x2 dilations, 2x2 strides
+  set_op({{1, 1, 2, 1}}, {{1, 2, 2, 1}}, "VALID", "NHWC");
+  INFER_OK(op, "[1,5,5,1];[3,3,1,1]", "[d0_0,2,1,d1_3]");
+
+  // Tests for NCHW
+  // 1x1 filter, 2x1 dilations, 1x1 strides
+  set_op({{1, 1, 2, 1}}, {{1, 1, 1, 1}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,2,2];[1,1,1,1]", "[d0_0,d1_3,2,2]");
+
+  // 1x1 filter, 2x1 dilations, 2x1 strides
+  set_op({{1, 1, 2, 1}}, {{1, 1, 2, 1}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,4,4];[1,1,1,1]", "[d0_0,d1_3,2,4]");
+
+  // 1x1 filter, 2x1 dilations, 2x2 strides
+  set_op({{1, 1, 2, 1}}, {{1, 1, 2, 2}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,4,4];[1,1,1,1]", "[d0_0,d1_3,2,2]");
+
+  // 3x3 filter, 2x1 dilations, 1x1 strides
+  set_op({{1, 1, 2, 1}}, {{1, 1, 1, 1}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,5,5];[3,3,1,1]", "[d0_0,d1_3,1,3]");
+
+  // 3x3 filter, 2x1 dilations, 2x1 strides
+  set_op({{1, 1, 2, 1}}, {{1, 1, 2, 1}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,5,5];[3,3,1,1]", "[d0_0,d1_3,1,3]");
+
+  // 3x3 filter, 1x2 dilations, 2x2 strides
+  set_op({{1, 1, 1, 2}}, {{1, 1, 2, 2}}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,5,5];[3,3,1,1]", "[d0_0,d1_3,2,1]");
+
+  // Some tests for "SAME" padding
+
+  // 4x4 input, 1x1 filter, 2x1 dilations, 1x1 stride
+  set_op({{1, 2, 1, 1}}, {{1, 1, 1, 1}}, "SAME", "NHWC");
+  INFER_OK(op, "[1,4,4,1];[1,1,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // 3x3 input, 2x2 filter, 2x2 dilations, 1x1 stride
+  set_op({{1, 2, 2, 1}}, {{1, 1, 1, 1}}, "SAME", "NHWC");
+  INFER_OK(op, "[1,3,3,1];[2,2,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+
+  // 4x4 input, 2x2 filter, 1x2 dilations, 2x2 stride
+  set_op({{1, 1, 2, 1}}, {{1, 2, 2, 1}}, "SAME", "NHWC");
+  INFER_OK(op, "[1,4,4,1];[2,2,1,1]", "[d0_0,2,2,d1_3]");
+
+  // 4x4 input, 2x2 filter, 2x2 dilations, 1x1 stride
+  set_op({{1, 2, 2, 1}}, {{1, 1, 1, 1}}, "SAME", "NHWC");
+  INFER_OK(op, "[1,4,4,1];[2,2,1,1]", "[d0_0,d0_1,d0_2,d1_3]");
+}
+
+TEST(CommonShapeFnsTest, Conv3DShapeTest) {
+  ShapeInferenceTestOp op("Conv3D");
+  auto set_op = [&op](const std::vector<int32>& strides,
+                      const string& padding) {
+    TF_CHECK_OK(NodeDefBuilder("test", "Conv3D")
+                    .Input("input", 0, DT_FLOAT)
+                    .Input("filter", 0, DT_FLOAT)
+                    .Attr("strides", strides)
+                    .Attr("padding", padding)
+                    .Finalize(&op.node_def));
+  };
+
+  // 1x1x1 filter
+  set_op({{1, 1, 1, 1, 1}}, "VALID");
+  INFER_OK(op, "[1,2,2,2,1];[1,1,1,1,1]", "[d0_0,2,2,2,d1_4]");
+
+  // Invalid rank for input
+  INFER_ERROR("must be rank 5", op, "[4,4];[2,1,1,1]");
+  // Invalid rank for filter
+  INFER_ERROR("must be rank 5", op, "[1,4,4,1];[2,1,1]");
+
+  // unknown dims in the critical fields give partial inference.
+  INFER_OK(op, "[1,2,2,2,1];[1,1,1,1,1]", "[d0_0,2,2,2,d1_4]");
+  INFER_OK(op, "[1,?,2,2,1];[1,1,1,1,1]", "[d0_0,?,2,2,d1_4]");
+  INFER_OK(op, "[1,2,?,2,1];[1,1,1,1,1]", "[d0_0,2,?,2,d1_4]");
+  INFER_OK(op, "[1,2,2,?,1];[1,1,1,1,1]", "[d0_0,2,2,?,d1_4]");
+  INFER_OK(op, "[1,2,2,2,1];[?,1,1,1,1]", "[d0_0,?,2,2,d1_4]");
+  INFER_OK(op, "[1,2,2,2,1];[1,?,1,1,1]", "[d0_0,2,?,2,d1_4]");
+  INFER_OK(op, "[1,2,2,2,1];[1,1,?,1,1]", "[d0_0,2,2,?,d1_4]");
+  INFER_OK(op, "[1,2,2,2,1];[1,1,1,?,1]", "[d0_0,2,2,2,d1_4]");
+  INFER_OK(op, "[1,2,2,2,1];[1,1,1,1,?]", "[d0_0,2,2,2,d1_4]");
+
+  // input depths must match.
+  INFER_ERROR("Dimensions must be equal, but are 10 and 10000", op,
+              "[1,2,2,2,10];[1,1,1,10000,20]");
+
+  // 2x2x2 filter
+  set_op({{1, 1, 1, 1, 1}}, "VALID");
+  INFER_OK(op, "[1,2,2,2,1];[2,2,2,1,1]", "[d0_0,1,1,1,d1_4]");
+
+  // 3x3 input, 1x1 filter, 2x2 stride
+  set_op({{1, 2, 2, 2, 1}}, "VALID");
+  INFER_OK(op, "[1,3,3,3,1];[1,1,1,1,1]", "[d0_0,2,2,2,d1_4]");
+
+  // 3x3 input, 1x1 filter, 2x1x1 stride
+  set_op({{1, 2, 1, 1, 1}}, "VALID");
+  INFER_OK(op, "[1,3,3,3,1];[1,1,1,1,1]", "[d0_0,2,3,3,d1_4]");
+
+  // 4x4 input, 2x2 filter, 1x1 stride
+  set_op({{1, 1, 1, 1, 1}}, "SAME");
+  INFER_OK(op, "[1,4,4,4,1];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+
+  // with SAME, filter doesn't matter except for last dim.
+  set_op({{1, 1, 1, 1, 1}}, "SAME");
+  INFER_OK(op, "[?,4,4,4,1];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,?,4,4,1];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,?,4,1];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,?,1];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,?];[2,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,1];[?,2,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,1];[2,?,2,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,1];[2,2,?,1,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,1];[2,2,2,?,1]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+  INFER_OK(op, "[1,4,4,4,1];[2,2,2,1,?]", "[d0_0,d0_1,d0_2,d0_3,d1_4]");
+
+  // with SAME, and stride != 1, division happens to produce output.
+  set_op({{1, 2, 3, 4, 1}}, "SAME");
+  INFER_OK(op, "[1,4,9,4,1];[2,2,2,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[?,4,9,4,1];[2,2,2,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,?,9,4,1];[2,2,2,1,1]", "[d0_0,?,3,1,d1_4]");
+  INFER_OK(op, "[1,4,?,4,1];[2,2,2,1,1]", "[d0_0,2,?,1,d1_4]");
+  INFER_OK(op, "[1,4,9,?,1];[2,2,2,1,1]", "[d0_0,2,3,?,d1_4]");
+  INFER_OK(op, "[1,4,9,4,?];[2,2,2,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,4,9,4,1];[?,2,2,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,4,9,4,1];[2,?,2,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,4,9,4,1];[2,2,?,1,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,4,9,4,1];[2,2,2,?,1]", "[d0_0,2,3,1,d1_4]");
+  INFER_OK(op, "[1,4,9,4,1];[2,2,2,1,?]", "[d0_0,2,3,1,d1_4]");
+}
+
+TEST(CommonShapeFnsTest, DepthwiseConv2DShapeTest) {
+  ShapeInferenceTestOp op("DepthwiseConv2dNative");
+  std::vector<int32> strides = {{1, 1, 1, 1}};
+  TF_CHECK_OK(NodeDefBuilder("test", "DepthwiseConv2dNative")
+                  .Input("input", 0, DT_FLOAT)
+                  .Input("filter", 0, DT_FLOAT)
+                  .Attr("strides", strides)
+                  .Attr("padding", "VALID")
+                  .Attr("data_format", "NHWC")
+                  .Finalize(&op.node_def));
+
+  // Most of DepthwiseConv2D is implicitly tested by Conv2D, so
+  // we test only the very-specific differences here.
+
+  // 1x1 filter, depth multiplication
+  INFER_OK(op, "[1,2,2,3];[1,1,3,4]", "[d0_0,2,2,12]");
+
+  // Input depths not compatible
+  INFER_ERROR("Dimensions must be equal, but are 3 and 12", op,
+              "[1,2,2,3];[1,1,12,4]");
+
+  // No unknown dims in the critical fields.
+  INFER_OK(op, "[1,2,2,3];[1,1,3,4]", "[d0_0,2,2,12]");
+  INFER_OK(op, "[1,?,2,3];[1,1,3,4]", "[d0_0,?,2,12]");
+  INFER_OK(op, "[1,2,?,3];[1,1,3,4]", "[d0_0,2,?,12]");
+  INFER_OK(op, "[1,2,2,3];[?,1,3,4]", "[d0_0,?,2,12]");
+  INFER_OK(op, "[1,2,2,3];[1,?,3,4]", "[d0_0,2,?,12]");
+  INFER_OK(op, "[1,2,2,3];[1,1,?,4]", "[d0_0,2,2,12]");
+  INFER_OK(op, "[1,2,2,?];[1,1,?,4]", "[d0_0,2,2,?]");
+  INFER_OK(op, "[1,2,2,3];[1,1,3,?]", "[d0_0,2,2,?]");
+
+  // Test for NCHW format.
+  TF_CHECK_OK(NodeDefBuilder("test", "DepthwiseConv2dNative")
+                  .Input("input", 0, DT_FLOAT)
+                  .Input("filter", 0, DT_FLOAT)
+                  .Attr("strides", strides)
+                  .Attr("padding", "VALID")
+                  .Attr("data_format", "NCHW")
+                  .Finalize(&op.node_def));
+
+  // 1x1 filter, depth multiplication
+  INFER_OK(op, "[1,3,2,2];[1,1,3,4]", "[d0_0,12,2,2]");
+}
+
+TEST(CommonShapeFnsTest, AvgPool2DShapeTest) {
+  ShapeInferenceTestOp op("AvgPool");
+  auto set_op = [&op](const std::vector<int32>& strides,
+                      const std::vector<int32>& ksizes, const string& padding,
+                      const string& data_format) {
+    TF_CHECK_OK(NodeDefBuilder("test", "AvgPool")
+                    .Input("input", 0, DT_FLOAT)
+                    .Attr("strides", strides)
+                    .Attr("ksize", ksizes)
+                    .Attr("padding", padding)
+                    .Attr("data_format", data_format)
+                    .Finalize(&op.node_def));
+  };
+
+  // Most of the functionality is tested by conv-like shapes,
+  // so we check the very-specific avgpooling features here.
+
+  // 1x1 filter, 1x1 stride
+  set_op({1, 1, 1, 1}, {1, 1, 1, 1}, "VALID", "NHWC");
+  INFER_OK(op, "[1,2,2,1]", "[d0_0,2,2,d0_3]");
+
+  // 4x4 input, 2x1 ksize, 1x2 stride
+  set_op({1, 1, 2, 1}, {1, 2, 1, 1}, "VALID", "NHWC");
+  INFER_OK(op, "[1,4,4,1]", "[d0_0,3,2,d0_3]");
+
+  // 4x4 input, 2x1 ksize, 1x2 stride
+  // unknown dims in the critical fields lead to partial inference.
+  // Assumes NHWC format.
+  INFER_OK(op, "[1,?,4,1]", "[d0_0,?,2,d0_3]");
+  INFER_OK(op, "[1,4,?,1]", "[d0_0,3,?,d0_3]");
+
+  // 4x4 input, 2x1 ksize, 1x2 stride, NCHW format
+  set_op({{1, 1, 1, 2}}, {1, 1, 2, 1}, "VALID", "NCHW");
+  INFER_OK(op, "[1,1,4,4]", "[d0_0,d0_1,3,2]");
+
+  // 5x7 input, 2x2 ksize, 1x1 stride, NCHW_VECT_C test
+  set_op({{1, 1, 1, 1}}, {1, 1, 2, 2}, "VALID", "NCHW_VECT_C");
+  INFER_OK(op, "[2,3,5,7,4]", "[d0_0,d0_1,4,6,4]");
+  INFER_OK(op, "[5,7,?,?,4]", "[d0_0,d0_1,?,?,4]");
+  INFER_OK(op, "[?,?,?,?,4]", "[d0_0,d0_1,?,?,4]");
+  INFER_ERROR("Dimension must be 4 but is 3", op, "[2,5,7,11,3]");
+
+  // Invalid rank for input
+  INFER_ERROR("Shape must be rank", op, "[4,4]");
+}
+
+TEST(CommonShapeFnsTest, MaxPool2DShapeTest) {
+  ShapeInferenceTestOp op("MaxPool");
+  auto set_op = [&op](const std::vector<int32>& strides,
+                      const std::vector<int32>& ksizes, const string& padding,
+                      const string& data_format) {
+    TF_CHECK_OK(NodeDefBuilder("test", "MaxPool")
+                    .Input("input", 0, DT_FLOAT)
+                    .Attr("strides", strides)
+                    .Attr("ksize", ksizes)
+                    .Attr("padding", padding)
+                    .Attr("data_format", data_format)
+                    .Finalize(&op.node_def));
+  };
+
+  // Most of the functionality is tested by conv-like shapes,
+  // so we check the very-specific maxpooling features here,
+  // namely depthwise kernel and striding.
+
+  // all 1 strides, depth 2 filter
+  set_op({1, 1, 1, 1}, {1, 1, 1, 2}, "VALID", "NHWC");
+  INFER_OK(op, "[1,2,2,2]", "[d0_0,2,2,1]");
+
+  // depth 3 stride, 1x1x1 filter, NCHW
+  set_op({1, 3, 1, 1}, {1, 1, 1, 1}, "VALID", "NCHW");
+  INFER_OK(op, "[1,7,5,5]", "[d0_0,3,5,5]");
+
+  // 5x7 input, 2x2 ksize, 1x1 stride, NCHW_VECT_C tests
+  set_op({{1, 1, 1, 1}}, {1, 1, 2, 2}, "SAME", "NCHW_VECT_C");
+  INFER_OK(op, "[2,3,5,7,4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_OK(op, "[5,7,?,?,4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_OK(op, "[?,?,?,?,4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_ERROR("Dimension must be 4 but is 8", op, "[2,3,5,7,8]");
+}
+
+TEST(CommonShapeFnsTest, MaxPoolV22DShapeTest) {
+  ShapeInferenceTestOp op("MaxPoolV2");
+  Tensor ksizes_tensor, strides_tensor;
+  auto set_op = [&op, &ksizes_tensor, &strides_tensor](
+                    const std::vector<int32>& strides,
+                    const std::vector<int32>& ksizes, const string& padding,
+                    const string& data_format) {
+    TF_CHECK_OK(NodeDefBuilder("test", "MaxPoolV2")
+                    .Input("input", 0, DT_FLOAT)
+                    .Input("ksize", 1, DT_INT32)
+                    .Input("strides", 2, DT_INT32)
+                    .Attr("padding", padding)
+                    .Attr("data_format", data_format)
+                    .Finalize(&op.node_def));
+    ksizes_tensor = test::AsTensor<int32>(ksizes);
+    op.input_tensors.resize(3);
+    op.input_tensors[0] = nullptr;
+    op.input_tensors[1] = &ksizes_tensor;
+    strides_tensor = test::AsTensor<int32>(strides);
+    op.input_tensors[2] = &strides_tensor;
+  };
+
+  // Most of the functionality is tested by conv-like shapes,
+  // so we check the very-specific maxpooling features here,
+  // namely depthwise kernel and striding.
+
+  // all 1 strides, depth 2 filter
+  set_op({1, 1, 1, 1}, {1, 1, 1, 2}, "VALID", "NHWC");
+  INFER_OK(op, "[1,2,2,2];[4];[4]", "[d0_0,2,2,1]");
+
+  // depth 3 stride, 1x1x1 filter, NCHW
+  set_op({1, 3, 1, 1}, {1, 1, 1, 1}, "VALID", "NCHW");
+  INFER_OK(op, "[1,7,5,5];[4];[4]", "[d0_0,3,5,5]");
+
+  // 5x7 input, 2x2 ksize, 1x1 stride, NCHW_VECT_C tests
+  set_op({{1, 1, 1, 1}}, {1, 1, 2, 2}, "SAME", "NCHW_VECT_C");
+  INFER_OK(op, "[2,3,5,7,4];[4];[4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_OK(op, "[5,7,?,?,4];[4];[4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_OK(op, "[?,?,?,?,4];[4];[4]", "[d0_0,d0_1,d0_2,d0_3,4]");
+  INFER_ERROR("Dimension must be 4 but is 8", op, "[2,3,5,7,8];[4];[4]");
+}
+
+TEST(CommonShapeFnsTest, Pool3DShapeTest) {
+  ShapeInferenceTestOp op("MaxPool3D");
+  auto set_op = [&op](const std::vector<int32>& strides,
+                      const std::vector<int32>& ksizes, const string& padding) {
+    TF_CHECK_OK(NodeDefBuilder("test", "MaxPool3D")
+                    .Input("input", 0, DT_FLOAT)
+                    .Attr("strides", strides)
+                    .Attr("ksize", ksizes)
+                    .Attr("padding", padding)
+                    .Finalize(&op.node_def));
+  };
+
+  // Most of the functionality is tested by conv-like shapes,
+  // so we check that we handle the extra dimension properly.
+
+  // 2x3x4 stride, 1x1x1 filter.
+  set_op({1, 2, 3, 4, 1}, {1, 1, 1, 1, 1}, "VALID");
+  INFER_OK(op, "[1,24,24,24,1]", "[d0_0,12,8,6,d0_4]");
+
+  // Test partially known dimensions
+  set_op({1, 1, 3, 4, 1}, {1, 1, 1, 1, 1}, "VALID");
+  INFER_OK(op, "[1,?,24,24,1]", "[d0_0,?,8,6,d0_4]");
+}
+
+TEST(CommonShapeFnsTest, UnknownShapeTest) {
+  {
+    // Single output
+    ShapeInferenceTestOp op("QueueDequeue");
+    TF_CHECK_OK(NodeDefBuilder("test", "QueueDequeue")
+                    .Input("handle", 0, DT_STRING_REF)
+                    .Attr("component_types", {DT_FLOAT})
+                    .Finalize(&op.node_def));
+    INFER_OK(op, "[1]", "?");
+  }
+
+  {
+    // Multiple outputs
+    ShapeInferenceTestOp op("QueueDequeue");
+    TF_CHECK_OK(NodeDefBuilder("test", "QueueDequeue")
+                    .Input("handle", 0, DT_STRING_REF)
+                    .Attr("component_types", {DT_FLOAT, DT_FLOAT, DT_STRING})
+                    .Finalize(&op.node_def));
+    INFER_OK(op, "[1]", "?;?;?");
+  }
+}
+
+TEST(CommonShapeFnsTest, Reduce_ShapeFn) {
+  ShapeInferenceTestOp op("Sum");
+  op.input_tensors.resize(2);
+
+  TF_ASSERT_OK(NodeDefBuilder("test", "Sum")
+                   .Input("input", 0, DT_FLOAT)
+                   .Input("reduction_indices", 1, DT_INT32)
+                   .Attr("keep_dims", false)
+                   .Finalize(&op.node_def));
+
+  // Reduction indices not available, so output is unknown.
+  INFER_OK(op, "[2,4,5];[2]", "?");
+  INFER_OK(op, "?;[2]", "?");
+
+  Tensor indices = test::AsTensor<int32>({1, 2});
+  op.input_tensors[1] = &indices;
+
+  // Reduction indices available
+  INFER_OK(op, "[2,4,5];[2]", "[d0_0]");
+
+  // Wrapped indices
+  indices = test::AsTensor<int32>({-1, -2});
+  op.input_tensors[1] = &indices;
+  INFER_OK(op, "[2,4,5];[2]", "[d0_0]");
+
+  // Scalar
+  indices = test::AsScalar<int32>(0);
+  op.input_tensors[1] = &indices;
+  INFER_OK(op, "[2,4,5];[]", "[d0_1,d0_2]");
+
+  indices = test::AsScalar<int32>(-4);
+  op.input_tensors[1] = &indices;
+  INFER_ERROR("Invalid reduction dimension", op, "[2,4,5];[]");
+
+  // Empty reduction indices
+  indices = test::AsTensor<int32>({});
+  op.input_tensors[1] = &indices;
+  INFER_OK(op, "[2,4,5];[0]", "[d0_0,d0_1,d0_2]");
+
+  // Keep dims = true
+  TF_ASSERT_OK(NodeDefBuilder("test", "Sum")
+                   .Input("input", 0, DT_FLOAT)
+                   .Input("reduction_indices", 1, DT_INT32)
+                   .Attr("keep_dims", true)
+                   .Finalize(&op.node_def));
+  indices = test::AsTensor<int32>({-1, -2});
+  op.input_tensors[1] = &indices;
+  INFER_OK(op, "[2,4,5];[2]", "[d0_0, 1, 1]");
+
+  // input rank is known, but reduction indices are not (with keep_dim=true).
+  // The output rank matches input rank (because of keep_dims=true).
+  op.input_tensors[1] = nullptr;
+  INFER_OK(op, "[?,?,?];?", "[?,?,?]");
+  INFER_OK(op, "[?,?,?];[2]", "[?,?,?]");
+
+  // Reduction indices with too many dimensions.
+  INFER_ERROR("must be at most rank 1 but is rank 2", op, "[?,?,?];[?,?]");
+  // With older graph-def version, this is allowed.
+  op.graph_def_version = 20;
+  INFER_OK(op, "[?,?,?];[?,?]", "[?,?,?]");
+  // And when the tensor is specified, it's still allowed.
+  op.input_tensors[1] = &indices;
+  indices = test::AsTensor<int32>({-1, -2}, TensorShape({2, 1}));
+  INFER_OK(op, "[2,4,5];[2,1]", "[d0_0, 1, 1]");
+  indices = test::AsTensor<int32>({-1, -2}, TensorShape({1, 2}));
+  INFER_OK(op, "[2,4,5];[1,2]", "[d0_0, 1, 1]");
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownShapes) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {Unknown(), Unknown(), Unknown()}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownDims) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({-1, -1}), S({-1}), S({-1})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_InvalidIndicesRank) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({-1}), S({-1}), S({-1})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  EXPECT_EQ(error::INVALID_ARGUMENT,
+            ValidateSparseTensor(&c, indices, values, shape).code());
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_InvalidNumElements) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, 3}), S({4}), S({3})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  EXPECT_EQ(error::INVALID_ARGUMENT,
+            ValidateSparseTensor(&c, indices, values, shape).code());
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_InvalidRank) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, 3}), S({5}), S({4})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  EXPECT_EQ(error::INVALID_ARGUMENT,
+            ValidateSparseTensor(&c, indices, values, shape).code());
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownNumIndexElements) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({-1, 3}), S({5}), S({3})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownNumValueElements) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, 3}), S({-1}), S({3})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownIndexRank) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, -1}), S({5}), S({3})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor_UnknownShapeRank) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, 3}), S({5}), S({-1})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+TEST(CommonShapeFnsTest, ValidateSparseTensor) {
+  NodeDef def;
+  InferenceContext c(TF_GRAPH_DEF_VERSION, &def, MakeOpDef(3, 1),
+                     {S({5, 3}), S({5}), S({3})}, {}, {}, {});
+  EXPECT_EQ(3, c.num_inputs());
+  EXPECT_EQ(1, c.num_outputs());
+
+  auto indices = c.input(0);
+  auto values = c.input(1);
+  auto shape = c.input(2);
+  TF_EXPECT_OK(ValidateSparseTensor(&c, indices, values, shape));
+}
+
+}  // namespace shape_inference
+}  // namespace tensorflow

control_flow.h

@@ -0,0 +1,58 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_CONTROL_FLOW_H_
+#define TENSORFLOW_FRAMEWORK_CONTROL_FLOW_H_
+
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+const uint64 kIllegalFrameId = ~0uLL;
+const int64 kIllegalIterId = -1;
+
+// For the purpose of control flow, every tensor produced by TensorFlow is
+// conceptually tagged by a 'FrameAndIter'. FrameAndIter consists of a
+// 'frame_id' and an 'iter_id'. The tensor value it represents is produced
+// in the frame with frame_id at the iteration of iter_id.
+struct FrameAndIter {
+  uint64 frame_id = kIllegalFrameId;
+  int64 iter_id = kIllegalIterId;
+
+  FrameAndIter() {}
+
+  FrameAndIter(uint64 frame, int64 iter) {
+    frame_id = frame;
+    iter_id = iter;
+  }
+
+  bool operator==(const FrameAndIter& other) const {
+    return (frame_id == other.frame_id && iter_id == other.iter_id);
+  }
+};
+
+struct FrameAndIterHash {
+  size_t operator()(const FrameAndIter& key) const {
+    // Make sure there are no padding bytes that we don't want
+    CHECK_EQ(sizeof(uint64) + sizeof(int64), sizeof(FrameAndIter));
+    return Hash64(reinterpret_cast<const char*>(&key), sizeof(FrameAndIter));
+  }
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_CONTROL_FLOW_H_

cost_graph.proto

@@ -0,0 +1,72 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "CostGraphProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
+
+message CostGraphDef {
+  message Node {
+    // The name of the node. Names are globally unique.
+    string name = 1;
+
+    // The device of the node. Can be empty if the node is mapped to the
+    // default partition or partitioning hasn't been run yet.
+    string device = 2;
+
+    // The id of the node. Node ids are only unique inside a partition.
+    int32 id = 3;
+
+    // Inputs of this node. They must be executed before this node can be
+    // executed. An input is a particular output of another node, specified
+    // by the node id and the output index.
+    message InputInfo {
+      int32 preceding_node = 1;
+      int32 preceding_port = 2;
+    }
+    repeated InputInfo input_info = 4;
+
+    // Outputs of this node.
+    message OutputInfo {
+      int64 size = 1;
+      // If >= 0, the output is an alias of an input. Note that an alias input
+      // may itself be an alias. The algorithm will therefore need to follow
+      // those pointers.
+      int64 alias_input_port = 2;
+      TensorShapeProto shape = 3;
+      DataType dtype = 4;
+    }
+    repeated OutputInfo output_info = 5;
+
+    // Temporary memory used by this node.
+    int64 temporary_memory_size = 6;
+
+    int64 host_temp_memory_size = 10;
+    int64 device_temp_memory_size = 11;
+    int64 host_persistent_memory_size = 12;
+    int64 device_persistent_memory_size = 16;
+
+    // Estimate of the computational cost of this node, in microseconds.
+    int64 compute_cost = 9;
+
+    // Analytical estimate of the computational cost of this node, in
+    // microseconds.
+    int64 compute_time = 14;
+
+    // Analytical estimate of the memory access cost of this node, in
+    // microseconds.
+    int64 memory_time = 15;
+
+    // If true, the output is permanent: it can't be discarded, because this
+    // node is part of the "final output". Nodes may depend on final nodes.
+    bool is_final = 7;
+
+    // Ids of the control inputs for this node.
+    repeated int32 control_input = 8;
+  }
+  repeated Node node = 1;
+}

device_attributes.proto

@@ -0,0 +1,35 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "DeviceAttributesProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+message DeviceLocality {
+  // Optional bus locality of device.  Default value of 0 means
+  // no specific locality.  Specific localities are indexed from 1.
+  int32 bus_id = 1;
+};
+
+message DeviceAttributes {
+  // Fully specified name of the device within a cluster.
+  string name = 1;
+
+  // String representation of device_type.
+  string device_type = 2;
+
+  // Memory capacity of device in bytes.
+  int64 memory_limit = 4;
+
+  // Platform-specific data about device that may be useful
+  // for supporting efficient data transfers.
+  DeviceLocality locality = 5;
+
+  // A device is assigned a global unique number each time it is
+  // initialized. "incarnation" should never be 0.
+  fixed64 incarnation = 6;
+
+  // String representation of the physical device that this device maps to.
+  string physical_device_desc = 7;
+}

device_base.cc

@@ -0,0 +1,30 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/device_base.h"
+
+namespace tensorflow {
+
+DeviceBase::~DeviceBase() {}
+
+const DeviceAttributes& DeviceBase::attributes() const {
+  LOG(FATAL) << "Device does not implement attributes()";
+}
+
+const string& DeviceBase::name() const {
+  LOG(FATAL) << "Device does not implement name()";
+}
+
+}  // namespace tensorflow

device_base.h

@@ -0,0 +1,243 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_
+#define TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/refcount.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace Eigen {
+struct ThreadPoolDevice;
+#ifdef TENSORFLOW_USE_SYCL
+struct SyclDevice;
+#endif
+}  // end namespace Eigen
+
+namespace perftools {
+namespace gputools {
+class Stream;
+}  // namespace gputools
+}  // namespace perftools
+
+namespace tensorflow {
+
+class Device;
+class DeviceAttributes;
+class Env;
+class EventMgr;
+class OpKernelContext;
+class ResourceMgr;
+class TensorProto;
+
+namespace thread {
+class ThreadPool;
+}
+
+// A wrapper for an Eigen Gpu Device that includes per-op state. The
+// class is defined even for non-GPU devices since the
+// OpKernelContext::Params structure wants to fill it in.
+class PerOpGpuDevice {
+ public:
+  virtual ~PerOpGpuDevice() {}
+  virtual const Eigen::GpuDevice& device() const = 0;
+};
+
+// A class that devices can subclass to pass around
+// Device-specific context to OpKernels.
+class DeviceContext : public core::RefCounted {
+ public:
+  ~DeviceContext() override {}
+  virtual perftools::gputools::Stream* stream() const { return nullptr; }
+  virtual void MaintainLifetimeOnStream(
+      const Tensor* t, perftools::gputools::Stream* stream) const {}
+
+  // "cpu_tensor" is a tensor on a CPU. Copies "cpu_tensor" into
+  // "device_tensor" which is on a GPU device "device". "device_tensor"
+  // must be allocated to be of the same size as "cpu_tensor".
+  virtual void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
+                                     Tensor* device_tensor,
+                                     StatusCallback done) const {
+    done(errors::Internal("Unrecognized device type in CPU-to-device Copy"));
+  }
+
+  // "device_tensor" is a tensor on a non-CPU device.  Copies
+  // device_tensor into "cpu_tensor".  "cpu_tensor" must be allocated
+  // to be of the same size as "device_tensor".
+  virtual void CopyDeviceTensorToCPU(const Tensor* device_tensor,
+                                     StringPiece tensor_name, Device* device,
+                                     Tensor* cpu_tensor, StatusCallback done) {
+    done(errors::Internal("Unrecognized device type in device-to-CPU Copy"));
+  }
+};
+
+// map[i] is the DeviceContext* for the node with id i, if i < map.size().
+typedef std::vector<DeviceContext*> DeviceContextMap;
+
+class DeviceBase {
+ public:
+  explicit DeviceBase(Env* env) : env_(env) {}
+  virtual ~DeviceBase();
+
+  Env* env() const { return env_; }
+
+  // Override this to return true for devices that require an Op's
+  // compute method to save references to the temporary tensors it
+  // allocates until the Op execution completes
+  virtual bool RequiresRecordingAccessedTensors() const { return false; }
+
+  struct CpuWorkerThreads {
+    int num_threads = 0;
+    thread::ThreadPool* workers = nullptr;
+  };
+
+  // Does not take ownership.
+  void set_tensorflow_cpu_worker_threads(CpuWorkerThreads* t) {
+    cpu_worker_threads_ = t;
+  }
+
+  virtual const CpuWorkerThreads* tensorflow_cpu_worker_threads() const {
+    CHECK(cpu_worker_threads_ != nullptr);
+    return cpu_worker_threads_;
+  }
+
+  // "stream" is used in special circumstances (such as the
+  // constructors of Ops) where there is no available OpKernelContext.
+  // "default_context" is used by OpKernelContext whenever a device does not
+  // supply a DeviceContext for an op in FillContextMap (e.g. when only
+  // using a single stream.)
+  // "event_mgr" is used to delay deallocation of temporary GPU buffers.
+  // TODO(pbar) Work out how to move this out of DeviceBase.
+  struct GpuDeviceInfo {
+    // Make sure all the defaults are NULL, so we can spot missing assignments.
+    perftools::gputools::Stream* stream = nullptr;
+    DeviceContext* default_context = nullptr;
+    EventMgr* event_mgr = nullptr;
+    int gpu_id = -1;
+  };
+
+  // Does not take ownership.
+  void set_tensorflow_gpu_device_info(GpuDeviceInfo* g) {
+    gpu_device_info_ = g;
+  }
+
+  virtual const GpuDeviceInfo* tensorflow_gpu_device_info() const {
+    return gpu_device_info_;
+  }
+
+  // The preferred thread pool for this device. If it is nullptr, the system
+  // automatically assigns a thread pool for execution.
+  virtual thread::ThreadPool* tensorflow_device_thread_pool() {
+    return device_thread_pool_;
+  }
+
+  // Does not take ownership.
+  void set_eigen_cpu_device(Eigen::ThreadPoolDevice* d) {
+    eigen_cpu_device_ = d;
+  }
+
+#ifdef TENSORFLOW_USE_SYCL
+  void set_eigen_sycl_device(Eigen::SyclDevice* d) { eigen_sycl_device_ = d; }
+#endif
+
+  // Return the Allocator implementation to use based on the allocator
+  // attributes requested.  See allocator.h for more details.
+  virtual Allocator* GetAllocator(AllocatorAttributes /*attr*/) {
+    LOG(FATAL) << "GetAllocator() is not implemented.";
+    return nullptr;
+  }
+
+  // Return the Allocator implementation to use based on the allocator
+  // attributes requested and the supplied resource manager. By
+  // default this ignores the resource manager and calls the base
+  // implementation but devices can override if they want to consult
+  // the resource manager when choosing the allocator.
+  virtual Allocator* GetStepAllocator(AllocatorAttributes attr,
+                                      ResourceMgr* /*step_resource_manager*/) {
+    return GetAllocator(attr);
+  }
+
+  virtual const Eigen::ThreadPoolDevice* eigen_cpu_device() {
+    CHECK(eigen_cpu_device_ != nullptr);
+    return eigen_cpu_device_;
+  }
+
+#ifdef TENSORFLOW_USE_SYCL
+  virtual const Eigen::SyclDevice* eigen_sycl_device() const {
+    CHECK(eigen_sycl_device_ != nullptr);
+    return eigen_sycl_device_;
+  }
+#endif
+
+  // Caller owns the return value. The OpKernelContext calls this even
+  // for devices that do not implement an eigen_gpu_device. Overridden
+  // by GPU devices to return a derived type.
+  virtual PerOpGpuDevice* MakeGpuDevice() { return nullptr; }
+
+  virtual DeviceBase* UnderlyingDevice() { return this; }
+  virtual const DeviceBase* UnderlyingDevice() const { return this; }
+
+  // This is overridden by GPU devices to reinitialize the derived
+  // type returned by MakeGpuDevice.
+  virtual void ReinitializeGpuDevice(OpKernelContext* /*context*/,
+                                     PerOpGpuDevice* /*device*/,
+                                     DeviceContext* /*dc*/,
+                                     Allocator* /*allocator*/) {}
+
+  // Unimplemented by default
+  virtual const DeviceAttributes& attributes() const;
+  virtual const string& name() const;
+
+  // Materializes the given TensorProto into 'tensor' stored in Device
+  // memory.  Most devices will want to override this.
+  //
+  // TODO(vrv): We should be able to put this function into
+  // OpKernelContext and handle the copies from device memory via send
+  // and receive nodes, instead of requiring that each device handle
+  // the copies here as well as in copy ops.
+  virtual Status MakeTensorFromProto(const TensorProto& tensor_proto,
+                                     const AllocatorAttributes alloc_attrs,
+                                     Tensor* tensor) {
+    return errors::Internal("Device does not implement MakeTensorFromProto()");
+  }
+
+ protected:
+  // Does not take ownership.
+  void set_tensorflow_device_thread_pool(thread::ThreadPool* thread_pool) {
+    device_thread_pool_ = thread_pool;
+  }
+
+ private:
+  Env* const env_;
+  CpuWorkerThreads* cpu_worker_threads_ = nullptr;
+  GpuDeviceInfo* gpu_device_info_ = nullptr;
+  thread::ThreadPool* device_thread_pool_ = nullptr;
+  Eigen::ThreadPoolDevice* eigen_cpu_device_ = nullptr;
+#ifdef TENSORFLOW_USE_SYCL
+  Eigen::SyclDevice* eigen_sycl_device_ = nullptr;
+#endif
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_

fake_input.cc

@@ -0,0 +1,240 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/fake_input.h"
+
+#include <vector>
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_def_util.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace {
+
+class FakeInputImpl {
+ public:
+  FakeInputImpl(const OpDef* op_def, int in_index, const NodeDef* node_def,
+                NodeDefBuilder* builder);
+  void SetN(int n);
+  void SetDataType(DataType dt);
+  void SetTypeList(DataTypeSlice dts);
+  Status AddInputToBuilder();
+
+ private:
+  static string FakeNodeName(int in_index);
+  Status GetN(int* n) const;
+  Status GetDataType(DataType* dt) const;
+  void NSources(int n, DataType dt) const;
+  void SourceList(DataTypeSlice dts) const;
+
+  const OpDef* const op_def_;
+  const OpDef::ArgDef* const arg_;
+  const string in_node_;
+  const NodeDef* const node_def_;
+  NodeDefBuilder* const builder_;
+
+  bool n_specified_;
+  int n_;
+  bool dt_specified_;
+  DataType dt_;
+  bool dts_specified_;
+  DataTypeSlice dts_;
+};
+
+FakeInputImpl::FakeInputImpl(const OpDef* op_def, int in_index,
+                             const NodeDef* node_def, NodeDefBuilder* builder)
+    : op_def_(op_def),
+      arg_(&op_def->input_arg(in_index)),
+      in_node_(FakeNodeName(in_index)),
+      node_def_(node_def),
+      builder_(builder),
+      n_specified_(false),
+      dt_specified_(false),
+      dts_specified_(false) {}
+
+void FakeInputImpl::SetN(int n) {
+  n_specified_ = true;
+  n_ = n;
+}
+
+void FakeInputImpl::SetDataType(DataType dt) {
+  dt_specified_ = true;
+  dt_ = dt;
+}
+
+void FakeInputImpl::SetTypeList(DataTypeSlice dts) {
+  dts_specified_ = true;
+  dts_ = dts;
+}
+
+Status FakeInputImpl::AddInputToBuilder() {
+  if (dts_specified_) {
+    SourceList(dts_);
+
+  } else if (n_specified_ || !arg_->number_attr().empty()) {
+    int n;
+    TF_RETURN_IF_ERROR(GetN(&n));
+
+    DataType dt;
+    if (n > 0) {
+      TF_RETURN_IF_ERROR(GetDataType(&dt));
+    } else {
+      dt = DT_FLOAT;
+    }
+
+    NSources(n, dt);
+  } else {
+    if (!dt_specified_ && !arg_->type_list_attr().empty()) {
+      DataTypeVector dts;
+      Status status = GetNodeAttr(*node_def_, arg_->type_list_attr(), &dts);
+      if (!status.ok()) {
+        return errors::InvalidArgument(
+            "Could not infer list of types for input '", arg_->name(), "': ",
+            status.error_message());
+      }
+      SourceList(dts);
+      return Status::OK();
+    }
+
+    DataType dt;
+    TF_RETURN_IF_ERROR(GetDataType(&dt));
+    builder_->Input(in_node_, 0, dt);
+  }
+  return Status::OK();
+}
+
+// static
+string FakeInputImpl::FakeNodeName(int in_index) {
+  char c = 'a' + (in_index % 26);
+  return string(&c, 1);
+}
+
+Status FakeInputImpl::GetN(int* n) const {
+  if (n_specified_) {
+    *n = n_;
+  } else {
+    Status status = GetNodeAttr(*node_def_, arg_->number_attr(), n);
+    if (!status.ok()) {
+      return errors::InvalidArgument("Could not infer length of input '",
+                                     arg_->name(), "': ",
+                                     status.error_message());
+    }
+  }
+  return Status::OK();
+}
+
+Status FakeInputImpl::GetDataType(DataType* dt) const {
+  if (dt_specified_) {
+    *dt = dt_;
+    return Status::OK();  // Ignore is_ref field of arg_.
+  } else if (arg_->type() != DT_INVALID) {
+    *dt = arg_->type();
+  } else if (!arg_->type_attr().empty()) {
+    Status status = GetNodeAttr(*node_def_, arg_->type_attr(), dt);
+    if (!status.ok()) {
+      // Check if the type attr has a default
+      const OpDef::AttrDef* attr = FindAttr(arg_->type_attr(), *op_def_);
+      if (attr && attr->has_default_value()) {
+        *dt = attr->default_value().type();
+      } else {
+        return errors::InvalidArgument("Could not infer type for input '",
+                                       arg_->name(), "': ",
+                                       status.error_message());
+      }
+    }
+  } else {
+    return errors::InvalidArgument("No type or type_attr field in arg '",
+                                   arg_->name(), "'");
+  }
+  if (arg_->is_ref()) {
+    *dt = MakeRefType(*dt);
+  }
+  return Status::OK();
+}
+
+void FakeInputImpl::NSources(int n, DataType dt) const {
+  std::vector<NodeDefBuilder::NodeOut> srcs;
+  srcs.reserve(n);
+  for (int i = 0; i < n; ++i) {
+    srcs.emplace_back(in_node_, i, dt);
+  }
+  builder_->Input(gtl::ArraySlice<NodeDefBuilder::NodeOut>(srcs));
+}
+
+void FakeInputImpl::SourceList(DataTypeSlice dts) const {
+  std::vector<NodeDefBuilder::NodeOut> srcs;
+  srcs.reserve(dts.size());
+  for (size_t i = 0; i < dts.size(); ++i) {
+    srcs.emplace_back(in_node_, i, dts[i]);
+  }
+  builder_->Input(gtl::ArraySlice<NodeDefBuilder::NodeOut>(srcs));
+}
+
+}  // namespace
+
+// Public interface ------------------------------------------------------------
+
+FakeInputFunctor FakeInput() {
+  return [](const OpDef& op_def, int in_index, const NodeDef& node_def,
+            NodeDefBuilder* builder) {
+    FakeInputImpl impl(&op_def, in_index, &node_def, builder);
+    return impl.AddInputToBuilder();
+  };
+}
+
+FakeInputFunctor FakeInput(DataType dt) {
+  return [dt](const OpDef& op_def, int in_index, const NodeDef& node_def,
+              NodeDefBuilder* builder) {
+    FakeInputImpl impl(&op_def, in_index, &node_def, builder);
+    impl.SetDataType(dt);
+    return impl.AddInputToBuilder();
+  };
+}
+
+FakeInputFunctor FakeInput(int n) {
+  return [n](const OpDef& op_def, int in_index, const NodeDef& node_def,
+             NodeDefBuilder* builder) {
+    FakeInputImpl impl(&op_def, in_index, &node_def, builder);
+    impl.SetN(n);
+    return impl.AddInputToBuilder();
+  };
+}
+
+FakeInputFunctor FakeInput(int n, DataType dt) {
+  return [n, dt](const OpDef& op_def, int in_index, const NodeDef& node_def,
+                 NodeDefBuilder* builder) {
+    FakeInputImpl impl(&op_def, in_index, &node_def, builder);
+    impl.SetN(n);
+    impl.SetDataType(dt);
+    return impl.AddInputToBuilder();
+  };
+}
+
+FakeInputFunctor FakeInput(DataTypeSlice dts) {
+  // Make a copy to ensure the data will still be around when the lambda is
+  // called.
+  DataTypeVector dtv(dts.begin(), dts.end());
+  return [dtv](const OpDef& op_def, int in_index, const NodeDef& node_def,
+               NodeDefBuilder* builder) {
+    FakeInputImpl impl(&op_def, in_index, &node_def, builder);
+    impl.SetTypeList(dtv);
+    return impl.AddInputToBuilder();
+  };
+}
+
+}  // namespace tensorflow

fake_input.h

@@ -0,0 +1,40 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_FAKE_INPUT_H_
+#define TENSORFLOW_FRAMEWORK_FAKE_INPUT_H_
+
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+
+// These functions return values that may be passed to
+// NodeDefBuilder::Input() to add an input for a test.  Use them when
+// you don't care about the node names/output indices providing the
+// input.  They also allow you to omit the input types and/or
+// list length when they may be inferred.
+FakeInputFunctor FakeInput();  // Infer everything
+FakeInputFunctor FakeInput(DataType dt);
+FakeInputFunctor FakeInput(int n);  // List of length n
+FakeInputFunctor FakeInput(int n, DataType dt);
+FakeInputFunctor FakeInput(DataTypeSlice dts);
+inline FakeInputFunctor FakeInput(std::initializer_list<DataType> dts) {
+  return FakeInput(DataTypeSlice(dts));
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_FAKE_INPUT_H_

function.cc

@@ -0,0 +1,1322 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/function.h"
+
+#include <map>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/function.pb_text.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/util/equal_graph_def.h"
+
+namespace tensorflow {
+
+// Extracts the actual type from "attr_values" based on its definition
+// "arg_def".
+//
+// If "arg_def" is a N*T type, *is_type_list is set to false, and
+// *dtypes is set to be a vector of size N and each element is T.
+//
+// If "arg_def" is a list(type), *is_type_list is set to true, and
+// *dtypes is set to be a vector of types specified in attrs for
+// arg_def.
+//
+// Otherwise (arg_def is a simple type T), *is_type_list is set to
+// false, and *dtypes is set to a single element vector, whose only
+// element is T.
+Status ArgNumType(AttrSlice attrs, const OpDef::ArgDef& arg_def,
+                  bool* is_type_list, DataTypeVector* dtypes) {
+  dtypes->clear();
+  if (!arg_def.type_list_attr().empty()) {
+    const AttrValue* v = attrs.Find(arg_def.type_list_attr());
+    if (v == nullptr) {
+      return errors::NotFound("type attr not found: ",
+                              arg_def.type_list_attr());
+    }
+    *is_type_list = true;
+    for (int i = 0; i < v->list().type_size(); ++i) {
+      dtypes->push_back(v->list().type(i));
+    }
+    return Status::OK();
+  }
+
+  *is_type_list = false;
+  int num = 1;
+  if (!arg_def.number_attr().empty()) {
+    const AttrValue* v = attrs.Find(arg_def.number_attr());
+    if (v == nullptr) {
+      return errors::NotFound("type attr not found: ", arg_def.type_attr());
+    }
+    num = v->i();
+  }
+
+  DataType dtype;
+  if (arg_def.type() != DT_INVALID) {
+    dtype = arg_def.type();
+  } else if (arg_def.type_attr().empty()) {
+    dtype = DT_INVALID;
+  } else {
+    const AttrValue* v = attrs.Find(arg_def.type_attr());
+    if (v == nullptr) {
+      return errors::NotFound("type attr not found: ", arg_def.type_attr());
+    }
+    dtype = v->type();
+  }
+  dtypes->resize(num, dtype);
+  return Status::OK();
+}
+
+namespace {
+
+template <typename T>
+void AddAttr(const string& name, const T& val, NodeDef* ndef) {
+  SetAttrValue(val, &((*ndef->mutable_attr())[name]));
+}
+
+Status ValidateSignatureWithAttrs(const OpDef& sig, AttrSlice attr_values) {
+  // attr_values should specify all attrs defined in fdef.
+  for (const auto& a : sig.attr()) {
+    const AttrValue* v = attr_values.Find(a.name());
+    if (!v) {
+      return errors::NotFound("Attr ", a.name(), " is not found from ",
+                              SummarizeOpDef(sig));
+    }
+    Status status = AttrValueHasType(*v, a.type());
+    if (!status.ok()) {
+      errors::AppendToMessage(&status, "for attr '", a.name(), "'");
+      return status;
+    }
+  }
+
+// TODO(josh11b): Enable this code once it works with function gradients.
+// Right now the C++ function gradient code assumes it can pass
+// all the attrs of the function to the gradient, and any attrs that
+// the gradient doesn't care about will be ignored.
+#if 0
+  if (attr_values.size() != sig.attr_size()) {
+    for (const auto& a : attr_values) {
+      // TODO(josh11b): Possibly should ignore attrs that start with "_" here?
+      bool found = false;
+      for (const auto& s : sig.attr()) {
+        if (a.first == s.name()) {
+          found = true;
+          break;
+        }
+      }
+      if (!found) {
+        return errors::NotFound("Attr ", a.first, " is not found in ",
+                                SummarizeOpDef(sig));
+      }
+    }
+  }
+#endif
+
+  return Status::OK();
+}
+
+// A helper class for instantiating functions. This contains shared information
+// like the resulting graph and node name index.
+class FunctionInstantiationHelper {
+ public:
+  FunctionInstantiationHelper(GetFunctionSignature get_function,
+                              InstantiationResult* result)
+      : get_function_(std ::move(get_function)), result_(*result) {
+    result_.nodes.clear();
+  }
+
+  // Builds index for nodes that can be used as node's input arguments.
+  Status BuildInputArgIndex(const OpDef::ArgDef& arg_def,
+                            AttrSlice attr_values) {
+    bool is_type_list;
+    DataTypeVector dtypes;
+    TF_RETURN_IF_ERROR(
+        ArgNumType(attr_values, arg_def, &is_type_list, &dtypes));
+    CHECK_GE(dtypes.size(), size_t{1});
+    int arg_index = result_.nodes.size();
+    TF_RETURN_IF_ERROR(
+        AddItem(arg_def.name(), {true, arg_index, 0, is_type_list, dtypes}));
+    // Creates dtypes.size() nodes in the graph.
+    for (size_t i = 0; i < dtypes.size(); ++i) {
+      TF_RETURN_IF_ERROR(AddItem(strings::StrCat(arg_def.name(), ":", i),
+                                 {true, arg_index, 0, false, {dtypes[i]}}));
+      DCHECK_EQ(arg_index, result_.nodes.size());
+      string name = arg_def.name();
+      if (dtypes.size() > 1) {
+        strings::StrAppend(&name, "_", i);
+      }
+      NodeDef* gnode = AddNode(name);
+      gnode->set_op("_Arg");
+      AddAttr("T", dtypes[i], gnode);
+      AddAttr("index", arg_index, gnode);
+      result_.arg_types.push_back(dtypes[i]);
+      ++arg_index;
+    }
+    return Status::OK();
+  }
+
+  Status BuildNodeOutputIndex(const NodeDef& node, AttrSlice attrs,
+                              const int arg_index) {
+    const OpDef* node_sig = nullptr;
+    TF_RETURN_IF_ERROR(get_function_(node.op(), &node_sig));
+    if (node_sig->output_arg_size() == 0) {
+      return AddItem(node.name(), {false, arg_index, 0, false, {}});
+    }
+    const int num_retval = node_sig->output_arg_size();
+    int start = 0;
+    bool is_type_list;
+    DataTypeVector dtypes;
+    for (int i = 0; i < num_retval; ++i) {
+      TF_RETURN_IF_ERROR(
+          ArgNumType(attrs, node_sig->output_arg(i), &is_type_list, &dtypes));
+      // Note that we rely on the backwards-compatibility test enforcing
+      // that output_arg(*).name() doesn't change here.
+      const string base_name =
+          strings::StrCat(node.name(), ":", node_sig->output_arg(i).name());
+      TF_RETURN_IF_ERROR(
+          AddItem(base_name, {false, arg_index, start, is_type_list, dtypes}));
+      for (int j = 0; j < static_cast<int>(dtypes.size()); ++j) {
+        TF_RETURN_IF_ERROR(
+            AddItem(strings::StrCat(base_name, ":", j),
+                    {false, arg_index, start + j, false, {dtypes[j]}}));
+      }
+      start += dtypes.size();
+    }
+    return Status::OK();
+  }
+
+  Status InstantiateNode(const NodeDef& fnode, AttrSlice attrs) {
+    const OpDef* fnode_sig = nullptr;
+    TF_CHECK_OK(get_function_(fnode.op(), &fnode_sig));
+    NodeDef* gnode = AddNode(fnode.name());
+    gnode->set_op(fnode.op());
+    gnode->set_device(fnode.device());
+    int gnode_idx = nodes_.size() - 1;
+
+    // Input
+    const int num_args = fnode_sig->input_arg_size();
+    bool is_type_list;  // ignored
+    DataTypeVector dtypes;
+    int fnode_arg_index = 0;
+    for (int i = 0; i < num_args; ++i) {
+      TF_RETURN_IF_ERROR(
+          ArgNumType(attrs, fnode_sig->input_arg(i), &is_type_list, &dtypes));
+      // Consume inputs (indexed by fnode_arg_index) until we have
+      // matched each element of dtypes (indexed by j).
+      for (size_t j = 0; j < dtypes.size(); ++fnode_arg_index) {
+        if (fnode_arg_index >= fnode.input_size()) {
+          // Should never happen if we computed dtypes correctly.
+          return errors::InvalidArgument(
+              "Attempt to access beyond input size: ", fnode_arg_index,
+              " >= ", fnode.input_size());
+        }
+        // Look up the next input.
+        const string& input_name = fnode.input(fnode_arg_index);
+        const auto* item = GetItemOrNull(input_name);
+        if (item == nullptr) {
+          return errors::InvalidArgument(
+              "input ", input_name, " is not found: ", SummarizeNodeDef(fnode));
+        }
+        if (item->dtypes.size() > dtypes.size() - j) {
+          return errors::InvalidArgument("Input ", input_name, " too long for ",
+                                         fnode_sig->input_arg(i).name());
+        }
+        // Match up all the elements of this input (indexed by k) with
+        // elements of dtypes (advancing j).
+        for (int k = 0; k < item->dtypes.size(); ++k, ++j) {
+          if (item->dtypes[k] != dtypes[j]) {
+            return errors::InvalidArgument(
+                "input ", fnode_sig->input_arg(i).name(), "[", j,
+                "] expected type ", DataTypeString(dtypes[j]),
+                " != ", DataTypeString(item->dtypes[k]), ", the type of ",
+                input_name, "[", k, "]");
+          }
+          if (item->is_func_arg) {
+            AddInput(gnode_idx, item->nid + k, 0);
+          } else {
+            AddInput(gnode_idx, item->nid, item->idx + k);
+          }
+        }
+      }
+    }
+
+    // Control deps.
+    for (int i = fnode_arg_index; i < fnode.input_size(); ++i) {
+      const string& input = fnode.input(i);
+      if (input.empty() || input[0] != '^') {
+        return errors::InvalidArgument("Expected input[", i, "] == '", input,
+                                       "' to be a control input.");
+      }
+      int nid = -1;
+      const string node_name = input.substr(1);
+      const string node_colon = node_name + ":";
+      const string node_colon_bound = node_name + ";";
+      // index_ is a map sorted lexicographically, so the key we are looking for
+      // must lie in the range [node_name, node_colon_bound).
+      auto it = index_.lower_bound(node_name);
+      while (it != index_.end() && it->first <= node_colon_bound) {
+        if (it->first == node_name ||
+            tensorflow::StringPiece(it->first).starts_with(node_colon)) {
+          nid = it->second.nid;
+          break;
+        }
+        ++it;
+      }
+      if (nid == -1) {
+        return errors::InvalidArgument("input[", i, "] == '", input,
+                                       "', is not found.");
+      }
+      AddDep(gnode_idx, nid);
+    }
+
+    // Attrs.
+    for (const auto& p : attrs) {
+      (*gnode->mutable_attr())[p.first] = p.second;
+    }
+
+    return Status::OK();
+  }
+
+  Status AddReturnNode(
+      const OpDef::ArgDef& ret_def, AttrSlice attrs,
+      const ::tensorflow::protobuf::Map<string, string>& ret_map,
+      int* ret_index) {
+    auto ret_iter = ret_map.find(ret_def.name());
+    if (ret_iter == ret_map.end()) {
+      return errors::InvalidArgument("Return ", ret_def.name(), " missing.");
+    }
+    bool is_type_list;
+    DataTypeVector dtypes;
+    TF_RETURN_IF_ERROR(ArgNumType(attrs, ret_def, &is_type_list, &dtypes));
+    CHECK_GE(dtypes.size(), size_t{1});
+    const auto* item = GetItemOrNull(ret_iter->second);
+    if (item == nullptr) {
+      return errors::InvalidArgument("Return ", ret_def.name(), " -> ",
+                                     ret_iter->second, " is not found.");
+    }
+    if (dtypes != item->dtypes) {
+      return errors::InvalidArgument("Invalid ret types ", ret_def.name(),
+                                     " : ", DataTypeVectorString(dtypes),
+                                     " vs. ",
+                                     DataTypeVectorString(item->dtypes));
+    }
+    for (size_t i = 0; i < dtypes.size(); ++i) {
+      string name = strings::StrCat(ret_def.name(), "_RetVal");
+      if (dtypes.size() > 1) {
+        strings::StrAppend(&name, "_", i);
+      }
+      NodeDef* gnode = AddNode(name);
+      gnode->set_op("_Retval");
+      AddInput(nodes_.size() - 1, item->nid, item->idx + i);
+      AddAttr("T", dtypes[i], gnode);
+      AddAttr("index", (*ret_index)++, gnode);
+      result_.ret_types.push_back(dtypes[i]);
+    }
+    return Status::OK();
+  }
+
+  // Adds the actual node inputs to the result graph by converting indexes to
+  // the node names.
+  void AddNodeInputs() {
+    for (int i = 0; i < result_.nodes.size(); i++) {
+      NodeInfo& node_info = nodes_[i];
+      for (const auto& p : node_info.data_inputs) {
+        result_.nodes[i].add_input(Name(p.first, p.second));
+      }
+      for (int index : node_info.control_inputs) {
+        result_.nodes[i].add_input(Dep(index));
+      }
+    }
+  }
+
+ private:
+  // This is used to build a small index for all names that can be used as a
+  // node's input arguments.
+  //
+  // If is_func_arg is true, the name is a function's argument.  In
+  // this case, the produced graph def has node[nid:nid + dtype.size()].
+  //
+  // Otherwise, the name is a function body's node return value.  In
+  // this case, the produced graph def has one node node[nid] and
+  // the node's output index [idx ... idx + num) corresponds to the
+  // named outputs.
+  //
+  // In all cases, "dtype" specifies the data type.
+  struct NameInfoItem {
+    bool is_func_arg;
+    int nid;
+    int idx;
+    bool is_type_list;
+    DataTypeVector dtypes;
+  };
+
+  // Adds an item into the input name index.
+  Status AddItem(const string& name, const NameInfoItem& item) {
+    if (!index_.insert({name, item}).second) {
+      return errors::InvalidArgument(
+          strings::StrCat("Duplicated ", item.is_func_arg ? "arg" : "ret",
+                          " name: "),
+          name);
+    }
+    return Status::OK();
+  }
+
+  const NameInfoItem* GetItemOrNull(const string& name) const {
+    return gtl::FindOrNull(index_, name);
+  }
+
+  string Dep(int node_index) const {
+    return strings::StrCat("^", Name(node_index));
+  }
+
+  string Name(int node_index) const {
+    CHECK_LT(node_index, nodes_.size());
+    return nodes_[node_index].name;
+  }
+
+  string Name(int node_index, int output_index) const {
+    if (output_index == 0) {
+      return Name(node_index);
+    } else {
+      return strings::StrCat(Name(node_index), ":", output_index);
+    }
+  }
+
+  NodeDef* AddNode(const string& name) {
+    result_.nodes.emplace_back();
+    NodeDef* gnode = &result_.nodes.back();
+    gnode->set_name(name);
+    nodes_.push_back({name, {}, {}});
+    CHECK_EQ(result_.nodes.size(), nodes_.size());
+    return gnode;
+  }
+
+  void AddInput(int node_index, int output_node, int output_index) {
+    CHECK_LT(node_index, nodes_.size());
+    nodes_[node_index].data_inputs.push_back(
+        std::make_pair(output_node, output_index));
+  }
+
+  void AddDep(int node_index, int dep_index) {
+    CHECK_LT(node_index, nodes_.size());
+    nodes_[node_index].control_inputs.push_back(dep_index);
+  }
+
+  GetFunctionSignature get_function_;
+  InstantiationResult& result_;
+  // A small index for all names that can be used as a node's input arguments.
+  std::map<string, NameInfoItem> index_;
+  // This contains information about a node in the new graph including the node
+  // names and input nodes' indexes.
+  struct NodeInfo {
+    string name;
+    // Data inputs where <n, k> means arg k of node n.
+    std::vector<std::pair<int, int>> data_inputs;
+    // Control inputs (dependencies).
+    std::vector<int> control_inputs;
+  };
+  // nodes_[i] is the information about result_.nodes[i].
+  std::vector<NodeInfo> nodes_;
+};
+
+// Various helpers Print(proto) to print relevant protos to ascii.
+string Print(const OpDef::ArgDef& arg) {
+  string out;
+  strings::StrAppend(&out, arg.name(), ":");
+  if (arg.is_ref()) strings::StrAppend(&out, "Ref(");
+  if (!arg.number_attr().empty()) {
+    strings::StrAppend(&out, arg.number_attr(), "*");
+  }
+  if (arg.type() != DT_INVALID) {
+    strings::StrAppend(&out, DataTypeString(arg.type()));
+  } else {
+    strings::StrAppend(&out, arg.type_attr());
+  }
+  if (arg.is_ref()) strings::StrAppend(&out, ")");
+  return out;
+}
+
+// TODO(josh11b): Merge this with SummarizeAttrValue().
+string Print(const AttrValue& attr_value) {
+  if (attr_value.value_case() == AttrValue::kType) {
+    return DataTypeString(attr_value.type());
+  } else if ((attr_value.value_case() == AttrValue::kList) &&
+             (attr_value.list().type_size() > 0)) {
+    string ret = "{";
+    for (int i = 0; i < attr_value.list().type_size(); ++i) {
+      if (i > 0) strings::StrAppend(&ret, ", ");
+      strings::StrAppend(&ret, DataTypeString(attr_value.list().type(i)));
+    }
+    strings::StrAppend(&ret, "}");
+    return ret;
+  } else if (attr_value.value_case() == AttrValue::kFunc) {
+    if (attr_value.func().attr_size() == 0) {
+      return attr_value.func().name();
+    }
+    std::vector<string> entries;
+    for (auto p : attr_value.func().attr()) {
+      entries.push_back(strings::StrCat(p.first, "=", Print(p.second)));
+    }
+    std::sort(entries.begin(), entries.end());
+    return strings::StrCat(attr_value.func().name(), "[",
+                           str_util::Join(entries, ", "), "]");
+  }
+  return SummarizeAttrValue(attr_value);
+}
+
+// TODO(josh11b): Merge this with SummarizeNodeDef().
+string Print(const NodeDef& n) {
+  string out;
+  strings::StrAppend(&out, n.name(), " = ", n.op());
+  if (n.attr_size() > 0) {
+    std::vector<string> entries;
+    for (auto& a : n.attr()) {
+      entries.push_back(strings::StrCat(a.first, "=", Print(a.second)));
+    }
+    std::sort(entries.begin(), entries.end());
+    strings::StrAppend(&out, "[", str_util::Join(entries, ", "), "]");
+  }
+  strings::StrAppend(&out, "(");
+  std::vector<StringPiece> dat;
+  std::vector<string> dep;
+  for (StringPiece s : n.input()) {
+    if (s.Consume("^")) {
+      dep.push_back(s.ToString());
+    } else {
+      dat.push_back(s);
+    }
+  }
+  strings::StrAppend(&out, str_util::Join(dat, ", "), ")");
+  if (!dep.empty()) {
+    strings::StrAppend(&out, " @ ", str_util::Join(dep, ", "));
+  }
+  return out;
+}
+
+string Print(const FunctionDef& fdef) {
+  string out;
+  const OpDef& sig = fdef.signature();
+  strings::StrAppend(&out, "\n", sig.name());
+  if (sig.attr_size() > 0) {
+    strings::StrAppend(&out, "[");
+    for (int i = 0; i < sig.attr_size(); ++i) {
+      const auto& a = sig.attr(i);
+      if (i > 0) strings::StrAppend(&out, ", ");
+      if (a.type() == "type") {
+        strings::StrAppend(&out, a.name(), ":", Print(a.allowed_values()));
+      } else {
+        strings::StrAppend(&out, a.name(), ":", a.type());
+      }
+    }
+    strings::StrAppend(&out, "]");
+  }
+  strings::StrAppend(&out, "(");
+  for (int i = 0; i < sig.input_arg_size(); ++i) {
+    if (i > 0) strings::StrAppend(&out, ", ");
+    strings::StrAppend(&out, Print(sig.input_arg(i)));
+  }
+  strings::StrAppend(&out, ") -> (");
+  for (int i = 0; i < sig.output_arg_size(); ++i) {
+    if (i > 0) strings::StrAppend(&out, ", ");
+    strings::StrAppend(&out, Print(sig.output_arg(i)));
+  }
+  strings::StrAppend(&out, ") {\n");
+  for (const auto& n : fdef.node_def()) {
+    strings::StrAppend(&out, "  ", Print(n), "\n");
+  }
+  for (const auto& r : fdef.ret()) {
+    strings::StrAppend(&out, "  return ", r.first, " = ", r.second, "\n");
+  }
+  strings::StrAppend(&out, "}\n");
+  return out;
+}
+
+string Print(gtl::ArraySlice<const NodeDef*> nodes) {
+  std::vector<const NodeDef*> arg;
+  std::vector<const NodeDef*> ret;
+  std::vector<const NodeDef*> body;
+  for (const NodeDef* n : nodes) {
+    if (n->op() == "_Arg") {
+      arg.push_back(n);
+    } else if (n->op() == "_Retval") {
+      ret.push_back(n);
+    } else {
+      body.push_back(n);
+    }
+  }
+  auto comp = [](const NodeDef* x, const NodeDef* y) {
+    int xi;
+    TF_CHECK_OK(GetNodeAttr(*x, "index", &xi));
+    int yi;
+    TF_CHECK_OK(GetNodeAttr(*y, "index", &yi));
+    return xi < yi;
+  };
+  std::sort(arg.begin(), arg.end(), comp);
+  std::sort(ret.begin(), ret.end(), comp);
+  string out;
+  strings::StrAppend(&out, "\n(");
+  auto get_type = [](const NodeDef& n) {
+    DataType dt;
+    if (!GetNodeAttr(n, "T", &dt).ok()) {
+      dt = DT_INVALID;
+    }
+    return DataTypeString(dt);
+  };
+  for (size_t i = 0; i < arg.size(); ++i) {
+    const NodeDef* n = arg[i];
+    if (i > 0) strings::StrAppend(&out, ", ");
+    CHECK_GE(n->attr_size(), 2);
+    strings::StrAppend(&out, n->name(), ":", get_type(*n));
+  }
+  strings::StrAppend(&out, ") -> (");
+  for (size_t i = 0; i < ret.size(); ++i) {
+    const NodeDef* n = ret[i];
+    if (i > 0) strings::StrAppend(&out, ", ");
+    CHECK_LE(2, n->attr_size());
+    CHECK_EQ(1, n->input_size());
+    strings::StrAppend(&out, n->input(0), ":", get_type(*n));
+  }
+  strings::StrAppend(&out, ") {\n");
+  for (size_t i = 0; i < body.size(); ++i) {
+    strings::StrAppend(&out, "  ", Print(*body[i]), "\n");
+  }
+  strings::StrAppend(&out, "}\n");
+  return out;
+}
+
+Status AddDefaultAttrs(const string& op,
+                       const GetFunctionSignature& get_function,
+                       AttrValueMap* attrs) {
+  const OpDef* op_def = nullptr;
+  TF_RETURN_IF_ERROR(get_function(op, &op_def));
+  AttrSlice attr_slice(attrs);
+  for (const auto& attr_def : op_def->attr()) {
+    if (attr_def.has_default_value() && !attr_slice.Find(attr_def.name())) {
+      if (!attrs->insert({attr_def.name(), attr_def.default_value()}).second) {
+        return errors::Internal("Somehow duplicated: ", attr_def.name());
+      }
+    }
+  }
+  return Status::OK();
+}
+
+}  // end namespace
+
+Status InstantiateFunction(const FunctionDef& fdef, AttrSlice attr_values,
+                           GetFunctionSignature get_function,
+                           InstantiationResult* result) {
+  VLOG(3) << "Instantiation Function: " << Print(fdef);
+
+  const OpDef& sig = fdef.signature();
+  TF_RETURN_IF_ERROR(ValidateSignatureWithAttrs(sig, attr_values));
+
+  FunctionInstantiationHelper helper(get_function, result);
+  Status s;
+  for (const OpDef::ArgDef& arg_def : sig.input_arg()) {
+    s = helper.BuildInputArgIndex(arg_def, attr_values);
+    if (!s.ok()) {
+      errors::AppendToMessage(&s, "In ", Print(arg_def));
+      return s;
+    }
+  }
+
+  auto substitute = [attr_values](StringPiece name, AttrValue* val) {
+    if (const AttrValue* v = attr_values.Find(name)) {
+      *val = *v;
+      return true;
+    }
+    return false;
+  };
+
+  // Makes a copy of all attrs in fdef and substitutes placeholders.
+  // After this step, every attr is bound to a concrete value.
+  std::vector<AttrValueMap> node_attrs;
+  node_attrs.resize(fdef.node_def_size());
+  for (int i = 0; i < fdef.node_def_size(); ++i) {
+    for (auto attr : fdef.node_def(i).attr()) {
+      if (!SubstitutePlaceholders(substitute, &attr.second)) {
+        return errors::InvalidArgument("Failed to bind all placeholders in ",
+                                       SummarizeAttrValue(attr.second));
+      }
+      if (!node_attrs[i].insert(attr).second) {
+        return errors::Internal("Somehow duplicated: ", attr.first);
+      }
+    }
+    TF_RETURN_IF_ERROR(
+        AddDefaultAttrs(fdef.node_def(i).op(), get_function, &node_attrs[i]));
+  }
+
+  for (int i = 0; i < fdef.node_def_size(); ++i) {
+    s = helper.BuildNodeOutputIndex(fdef.node_def(i), AttrSlice(&node_attrs[i]),
+                                    result->nodes.size() + i);
+    if (!s.ok()) {
+      errors::AppendToMessage(&s, "In ", SummarizeNodeDef(fdef.node_def(i)));
+      return s;
+    }
+  }
+  // Emits one node for each fdef.node_def.
+  for (int i = 0; i < fdef.node_def_size(); ++i) {
+    s = helper.InstantiateNode(fdef.node_def(i), AttrSlice(&node_attrs[i]));
+    if (!s.ok()) {
+      errors::AppendToMessage(&s, "In ", SummarizeNodeDef(fdef.node_def(i)));
+      return s;
+    }
+  }
+
+  // Emits nodes for the function's return values.
+  int ret_index = 0;
+  for (const OpDef::ArgDef& ret_def : sig.output_arg()) {
+    s = helper.AddReturnNode(ret_def, attr_values, fdef.ret(), &ret_index);
+    if (!s.ok()) {
+      errors::AppendToMessage(&s, "In function output ", Print(ret_def));
+      return s;
+    }
+  }
+
+  // Adds the actual node inputs using the input indexes.
+  helper.AddNodeInputs();
+
+  return Status::OK();
+}
+
+string DebugString(const FunctionDef& func_def) { return Print(func_def); }
+
+string DebugString(const GraphDef& instantiated_func_def) {
+  std::vector<const NodeDef*> ptrs;
+  for (const NodeDef& n : instantiated_func_def.node()) {
+    ptrs.push_back(&n);
+  }
+  return Print(ptrs);
+}
+
+string DebugString(gtl::ArraySlice<NodeDef> instantiated_func_nodes) {
+  std::vector<const NodeDef*> ptrs;
+  for (const NodeDef& n : instantiated_func_nodes) {
+    ptrs.push_back(&n);
+  }
+  return Print(ptrs);
+}
+
+string DebugStringWhole(const GraphDef& gdef) {
+  string ret;
+  for (const auto& fdef : gdef.library().function()) {
+    strings::StrAppend(&ret, Print(fdef));
+  }
+  strings::StrAppend(&ret, "\n");
+  for (const auto& ndef : gdef.node()) {
+    strings::StrAppend(&ret, Print(ndef), "\n");
+  }
+  return ret;
+}
+
+namespace {
+
+// Returns the name -> attr mapping of fdef's attrs that have a value set. In
+// Python, it's possible to access unset attrs, which returns a default value
+// and adds an unset attr to the map.
+std::map<string, AttrValue> GetSetAttrs(const FunctionDef& fdef) {
+  std::map<string, AttrValue> set_attrs;
+  for (auto pair : fdef.attr()) {
+    if (pair.second.value_case() != AttrValue::VALUE_NOT_SET) {
+      set_attrs[pair.first] = pair.second;
+    }
+  }
+  return set_attrs;
+}
+
+}  // end namespace
+
+bool FunctionDefsEqual(const FunctionDef& f1, const FunctionDef& f2) {
+  if (!OpDefEqual(f1.signature(), f2.signature())) return false;
+
+  std::map<string, AttrValue> f1_attrs = GetSetAttrs(f1);
+  std::map<string, AttrValue> f2_attrs = GetSetAttrs(f2);
+  if (f1_attrs.size() != f2_attrs.size()) return false;
+  for (auto iter1 : f1_attrs) {
+    auto iter2 = f2_attrs.find(iter1.first);
+    if (iter2 == f2_attrs.end()) return false;
+    if (!AreAttrValuesEqual(iter1.second, iter2->second)) return false;
+  }
+
+  if (!EqualRepeatedNodeDef(f1.node_def(), f2.node_def(), nullptr)) {
+    return false;
+  }
+
+  std::map<string, string> ret1(f1.ret().begin(), f1.ret().end());
+  std::map<string, string> ret2(f2.ret().begin(), f2.ret().end());
+  if (ret1 != ret2) return false;
+
+  return true;
+}
+
+uint64 FunctionDefHash(const FunctionDef& fdef) {
+  // signature
+  uint64 h = OpDefHash(fdef.signature());
+
+  // attrs
+  std::map<string, AttrValue> attrs = GetSetAttrs(fdef);
+  for (const auto& p : attrs) {
+    h = Hash64(p.first.data(), p.first.size(), h);
+    h = Hash64Combine(AttrValueHash(p.second), h);
+  }
+
+  // node defs
+  h = Hash64Combine(RepeatedNodeDefHash(fdef.node_def()), h);
+
+  // output names
+  std::map<string, string> ret(fdef.ret().begin(), fdef.ret().end());
+  for (const auto& p : ret) {
+    h = Hash64(p.first.data(), p.first.size(), h);
+    h = Hash64(p.second.data(), p.second.size(), h);
+  }
+
+  return h;
+}
+
+string Canonicalize(const string& funcname, AttrSlice attrs) {
+  std::vector<string> entries;
+  entries.reserve(attrs.size());
+  for (auto p : attrs) {
+    entries.push_back(strings::StrCat(p.first, "=", Print(p.second)));
+  }
+  std::sort(entries.begin(), entries.end());
+  return strings::StrCat(funcname, "[", str_util::Join(entries, ","), "]");
+}
+
+FunctionCallFrame::FunctionCallFrame(DataTypeSlice arg_types,
+                                     DataTypeSlice ret_types)
+    : arg_types_(arg_types.begin(), arg_types.end()),
+      ret_types_(ret_types.begin(), ret_types.end()) {
+  args_.resize(arg_types_.size());
+  rets_.resize(ret_types_.size());
+}
+
+FunctionCallFrame::~FunctionCallFrame() {}
+
+Status FunctionCallFrame::SetArgs(gtl::ArraySlice<Tensor> args) {
+  // Input type checks.
+  if (args.size() != arg_types_.size()) {
+    return errors::InvalidArgument("Expects ", arg_types_.size(),
+                                   " arguments, but ", args.size(),
+                                   " is provided");
+  }
+  for (size_t i = 0; i < args.size(); ++i) {
+    if (arg_types_[i] != args[i].dtype()) {
+      return errors::InvalidArgument(
+          "Expects arg[", i, "] to be ", DataTypeString(arg_types_[i]), " but ",
+          DataTypeString(args[i].dtype()), " is provided");
+    }
+    args_[i] = args[i];
+  }
+  return Status::OK();
+}
+
+Status FunctionCallFrame::GetRetvals(std::vector<Tensor>* rets) const {
+  rets->clear();
+  rets->reserve(rets_.size());
+  for (size_t i = 0; i < rets_.size(); ++i) {
+    const auto& item = rets_[i];
+    if (item.has_val) {
+      rets->push_back(item.val);
+    } else {
+      return errors::Internal("Retval[", i, "] does not have value");
+    }
+  }
+  return Status::OK();
+}
+
+Status FunctionCallFrame::ConsumeRetvals(std::vector<Tensor>* rets) {
+  rets->clear();
+  rets->reserve(rets_.size());
+  for (size_t i = 0; i < rets_.size(); ++i) {
+    if (rets_[i].has_val) {
+      rets->emplace_back(std::move(rets_[i].val));
+    } else {
+      return errors::Internal("Retval[", i, "] does not have value");
+    }
+  }
+  return Status::OK();
+}
+
+Status FunctionCallFrame::GetArg(int index, Tensor* val) const {
+  if (index < 0 || static_cast<size_t>(index) >= args_.size()) {
+    return errors::InvalidArgument("GetArg ", index, " is not within [0, ",
+                                   args_.size(), ")");
+  }
+  *val = args_[index];
+  return Status::OK();
+}
+
+Status FunctionCallFrame::SetRetval(int index, const Tensor& val) {
+  if (index < 0 || static_cast<size_t>(index) >= rets_.size()) {
+    return errors::InvalidArgument("SetRetval ", index, " is not within [0, ",
+                                   rets_.size(), ")");
+  }
+  if (val.dtype() != ret_types_[index]) {
+    return errors::InvalidArgument(
+        "Expects ret[", index, "] to be ", DataTypeString(ret_types_[index]),
+        ", but ", DataTypeString(val.dtype()), " is provided.");
+  }
+  Retval* item = &rets_[index];
+  if (!item->has_val) {
+    item->has_val = true;
+    item->val = val;
+  } else {
+    return errors::Internal("Retval[", index, "] has already been set.");
+  }
+  return Status::OK();
+}
+
+FunctionLibraryDefinition::FunctionDefAndOpRegistration::
+    FunctionDefAndOpRegistration(const FunctionDef& fdef_in)
+    : fdef(fdef_in),
+      // Exact shape inference for functions is handled by ShapeRefiner.
+      // Here we pass a dummy shape inference function for legacy code paths.
+      op_registration_data(fdef.signature(), shape_inference::UnknownShape,
+                           true /* is_function */) {}
+
+FunctionLibraryDefinition::FunctionLibraryDefinition(
+    const FunctionLibraryDefinition& other)
+    : default_registry_(other.default_registry_), func_grad_(other.func_grad_) {
+  for (const auto& it : other.function_defs_) {
+    TF_CHECK_OK(AddFunctionDef(it.second->fdef));
+  }
+}
+
+FunctionLibraryDefinition::FunctionLibraryDefinition(
+    const OpRegistryInterface* default_registry,
+    const FunctionDefLibrary& def_lib)
+    : default_registry_(default_registry),
+      function_defs_(def_lib.function_size()) {
+  for (const auto& fdef : def_lib.function()) {
+    // The latter function definition wins.
+    auto& ptr = function_defs_[fdef.signature().name()];
+    ptr.reset(new FunctionDefAndOpRegistration(fdef));
+  }
+  for (const auto& grad : def_lib.gradient()) {
+    func_grad_[grad.function_name()] = grad.gradient_func();
+  }
+}
+
+FunctionLibraryDefinition::~FunctionLibraryDefinition() {}
+
+const FunctionDef* FunctionLibraryDefinition::Find(const string& name) const {
+  auto iter = function_defs_.find(name);
+  if (iter == function_defs_.end()) {
+    return nullptr;
+  } else {
+    return &iter->second->fdef;
+  }
+}
+
+Status FunctionLibraryDefinition::AddFunctionDef(const FunctionDef& fdef) {
+  bool added;
+  return AddFunctionDefHelper(fdef, &added);
+}
+
+Status FunctionLibraryDefinition::AddFunctionDefHelper(const FunctionDef& fdef,
+                                                       bool* added) {
+  *added = false;
+  std::unique_ptr<FunctionDefAndOpRegistration>* entry =
+      &function_defs_[fdef.signature().name()];
+  if (*entry != nullptr) {
+    if (!FunctionDefsEqual((*entry)->fdef, fdef)) {
+      return errors::InvalidArgument(
+          "Cannot add function '", fdef.signature().name(),
+          "' because a different function with the same name already "
+          "exists.");
+    }
+    // Ignore duplicate FunctionDefs
+    return Status::OK();
+  }
+  const OpDef* op_def;
+  if (default_registry_->LookUpOpDef(fdef.signature().name(), &op_def).ok()) {
+    return errors::InvalidArgument(
+        "Cannot add function '", fdef.signature().name(),
+        "' because an op with the same name already exists.");
+  }
+  entry->reset(new FunctionDefAndOpRegistration(fdef));
+  *added = true;
+  return Status::OK();
+}
+
+Status FunctionLibraryDefinition::AddGradientDef(const GradientDef& grad) {
+  bool added;
+  return AddGradientDefHelper(grad, &added);
+}
+
+Status FunctionLibraryDefinition::AddGradientDefHelper(const GradientDef& grad,
+                                                       bool* added) {
+  *added = false;
+  string* entry = &func_grad_[grad.function_name()];
+  if (!entry->empty()) {
+    if (*entry != grad.gradient_func()) {
+      return errors::InvalidArgument(
+          "Cannot assign gradient function '", grad.gradient_func(), "' to '",
+          grad.function_name(), "' because it already has gradient function ",
+          "'", *entry, "'");
+    }
+    // Ignore duplicate GradientDefs
+    return Status::OK();
+  }
+  *entry = grad.gradient_func();
+  *added = true;
+  return Status::OK();
+}
+
+Status FunctionLibraryDefinition::AddLibrary(
+    const FunctionLibraryDefinition& other) {
+  // Remember the funcs and grads that we added successfully so that
+  // we can roll them back on error.
+  std::vector<string> funcs;
+  std::vector<string> funcs_with_grads;
+  Status s;
+  bool added;
+  for (auto iter : other.function_defs_) {
+    s = AddFunctionDefHelper(iter.second->fdef, &added);
+    if (!s.ok()) {
+      Remove(funcs, funcs_with_grads);
+      return s;
+    }
+    if (added) {
+      funcs.push_back(iter.second->fdef.signature().name());
+    }
+  }
+  for (auto iter : other.func_grad_) {
+    GradientDef grad;
+    grad.set_function_name(iter.first);
+    grad.set_gradient_func(iter.second);
+    s = AddGradientDefHelper(grad, &added);
+    if (!s.ok()) {
+      Remove(funcs, funcs_with_grads);
+      return s;
+    }
+    if (added) {
+      funcs_with_grads.push_back(grad.function_name());
+    }
+  }
+  return Status::OK();
+}
+
+Status FunctionLibraryDefinition::AddLibrary(
+    const FunctionDefLibrary& lib_def) {
+  // Remember the funcs and grads that we added successfully so that
+  // we can roll them back on error.
+  std::vector<string> funcs;
+  std::vector<string> funcs_with_grads;
+  Status s;
+  bool added;
+  for (const FunctionDef& fdef : lib_def.function()) {
+    s = AddFunctionDefHelper(fdef, &added);
+    if (!s.ok()) {
+      Remove(funcs, funcs_with_grads);
+      return s;
+    }
+    if (added) {
+      funcs.push_back(fdef.signature().name());
+    }
+  }
+  for (const GradientDef& grad : lib_def.gradient()) {
+    s = AddGradientDefHelper(grad, &added);
+    if (!s.ok()) {
+      Remove(funcs, funcs_with_grads);
+      return s;
+    }
+    if (added) {
+      funcs_with_grads.push_back(grad.function_name());
+    }
+  }
+  return Status::OK();
+}
+
+void FunctionLibraryDefinition::RemoveFunction(const string& func) {
+  const auto& i = function_defs_.find(func);
+  DCHECK(i != function_defs_.end());
+  function_defs_.erase(i);
+}
+
+void FunctionLibraryDefinition::RemoveGradient(const string& func) {
+  const auto& i = func_grad_.find(func);
+  DCHECK(i != func_grad_.end());
+  func_grad_.erase(i);
+}
+
+void FunctionLibraryDefinition::Remove(
+    const std::vector<string>& funcs,
+    const std::vector<string>& funcs_with_grads) {
+  for (const string& f : funcs) {
+    RemoveFunction(f);
+  }
+  for (const string& f : funcs_with_grads) {
+    RemoveGradient(f);
+  }
+}
+
+string FunctionLibraryDefinition::FindGradient(const string& func) const {
+  return gtl::FindWithDefault(func_grad_, func, "");
+}
+
+Status FunctionLibraryDefinition::LookUp(
+    const string& op, const OpRegistrationData** op_reg_data) const {
+  auto iter = function_defs_.find(op);
+  if (iter != function_defs_.end()) {
+    *op_reg_data = &iter->second->op_registration_data;
+    return Status::OK();
+  }
+  return default_registry_->LookUp(op, op_reg_data);
+}
+
+const FunctionDef* FunctionLibraryDefinition::GetAttrImpl(
+    const NodeDef& ndef) const {
+  if (ndef.op() != kGradientOp) {
+    // If 'ndef' calls a function and the function's def has the attr,
+    // returns it.
+    return Find(ndef.op());
+  }
+
+  // If ndef is SymbolicGradient[f=Foo], we use Foo's gradient or
+  // Foo's attributes.
+  const NameAttrList* forward_func_attrs;
+  if (!GetNodeAttr(ndef, kFuncAttr, &forward_func_attrs).ok()) {
+    return nullptr;
+  }
+  const string& func_name = forward_func_attrs->name();
+  const string& grad_name = FindGradient(func_name);
+  // If 'func' has a user-defined gradient function, uses the grad
+  // function's attrs to see if noinline is specified. Otherwise,
+  // uses func's attrs.
+  if (!grad_name.empty()) {
+    return Find(grad_name);
+  }
+  return Find(func_name);
+}
+
+FunctionDefLibrary FunctionLibraryDefinition::ToProto() const {
+  FunctionDefLibrary lib;
+  for (const auto& f : function_defs_) {
+    *lib.add_function() = f.second->fdef;
+  }
+  for (const auto& g : func_grad_) {
+    GradientDef* gd = lib.add_gradient();
+    gd->set_function_name(g.first);
+    gd->set_gradient_func(g.second);
+  }
+  return lib;
+}
+
+template <typename T>
+Status FunctionLibraryDefinition::GetAttr(const NodeDef& ndef,
+                                          const string& attr, T* value) const {
+  const FunctionDef* fdef = GetAttrImpl(ndef);
+  if (fdef && GetNodeAttr(AttrSlice(&fdef->attr()), attr, value).ok()) {
+    return Status::OK();
+  }
+  return errors::InvalidArgument("Attr ", attr, " is not defined.");
+}
+
+template <typename T>
+Status FunctionLibraryDefinition::GetAttr(const Node& node, const string& attr,
+                                          T* value) const {
+  return GetAttr(node.def(), attr, value);
+}
+
+#define GET_ATTR(T)                                                            \
+  template Status FunctionLibraryDefinition::GetAttr(const Node&,              \
+                                                     const string&, T*) const; \
+  template Status FunctionLibraryDefinition::GetAttr(const NodeDef&,           \
+                                                     const string&, T*) const;
+GET_ATTR(string)
+GET_ATTR(bool)
+#undef GET_ATTR
+
+void FunctionDefHelper::AttrValueWrapper::InitFromString(StringPiece val) {
+  if (val.size() >= 2 && val[0] == '$') {
+    proto.set_placeholder(val.data() + 1, val.size() - 1);
+  } else {
+    SetAttrValue(val, &proto);
+  }
+}
+
+FunctionDefHelper::AttrValueWrapper FunctionDefHelper::FunctionRef(
+    const string& name,
+    gtl::ArraySlice<std::pair<string, AttrValueWrapper>> attrs) {
+  AttrValueWrapper ret;
+  ret.proto.mutable_func()->set_name(name);
+  for (const auto& a : attrs) {
+    ret.proto.mutable_func()->mutable_attr()->insert({a.first, a.second.proto});
+  }
+  return ret;
+}
+
+NodeDef FunctionDefHelper::Node::ToNodeDef() const {
+  NodeDef n;
+  n.set_op(this->op);
+  n.set_name(this->ret[0]);
+  for (const auto& a : this->attr) {
+    n.mutable_attr()->insert({a.first, a.second.proto});
+  }
+  for (const string& a : this->arg) {
+    n.add_input(a);
+  }
+  for (const string& d : this->dep) {
+    n.add_input(strings::StrCat("^", d));
+  }
+  return n;
+}
+
+/* static */
+FunctionDef FunctionDefHelper::Create(
+    const string& function_name, gtl::ArraySlice<string> in_def,
+    gtl::ArraySlice<string> out_def, gtl::ArraySlice<string> attr_def,
+    gtl::ArraySlice<Node> node_def,
+    gtl::ArraySlice<std::pair<string, string>> ret_def) {
+  FunctionDef fdef;
+
+  // Signature
+  OpDefBuilder b(function_name);
+  for (const auto& i : in_def) b.Input(i);
+  for (const auto& o : out_def) b.Output(o);
+  for (const auto& a : attr_def) b.Attr(a);
+
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(b.Finalize(&op_reg_data));
+  fdef.mutable_signature()->Swap(&op_reg_data.op_def);
+
+  // Function body
+  for (const auto& n : node_def) {
+    *(fdef.add_node_def()) = n.ToNodeDef();
+  }
+
+  // Returns
+  for (const auto& r : ret_def) {
+    fdef.mutable_ret()->insert({r.first, r.second});
+  }
+  return fdef;
+}
+
+/* static */
+FunctionDef FunctionDefHelper::Define(const string& name,
+                                      gtl::ArraySlice<string> arg_def,
+                                      gtl::ArraySlice<string> ret_def,
+                                      gtl::ArraySlice<string> attr_def,
+                                      gtl::ArraySlice<Node> node_def) {
+  FunctionDef fdef;
+  OpDefBuilder b(name);
+  for (const auto& a : arg_def) b.Input(a);
+  for (const auto& r : ret_def) b.Output(r);
+  for (const auto& a : attr_def) b.Attr(a);
+
+  OpRegistrationData op_reg_data;
+  TF_CHECK_OK(b.Finalize(&op_reg_data));
+  fdef.mutable_signature()->Swap(&op_reg_data.op_def);
+
+  // Mapping from legacy output names to NodeDef outputs.
+  std::unordered_map<string, string> ret_index;
+  for (const auto& a : fdef.signature().input_arg()) {
+    ret_index[a.name()] = a.name();
+  }
+
+  // For looking up OpDefs
+  auto* op_def_registry = OpRegistry::Global();
+
+  // Function body
+  for (const auto& src : node_def) {
+    NodeDef* n = fdef.add_node_def();
+    n->set_op(src.op);
+    n->set_name(src.ret[0]);
+    for (const auto& a : src.attr) {
+      n->mutable_attr()->insert({a.first, a.second.proto});
+    }
+    for (const string& a : src.arg) {
+      const auto iter = ret_index.find(a);
+      CHECK(iter != ret_index.end()) << "Node input '" << a << "' in '"
+                                     << src.ret[0] << "' of " << name;
+      n->add_input(iter->second);
+    }
+    for (const string& d : src.dep) {
+      n->add_input(strings::StrCat("^", d));
+    }
+
+    // Add the outputs of this node to ret_index.
+    const OpDef* op_def = nullptr;
+    TF_CHECK_OK(op_def_registry->LookUpOpDef(n->op(), &op_def)) << n->op();
+    CHECK(op_def != nullptr) << n->op();
+    NameRangeMap output_names;
+    TF_CHECK_OK(NameRangesForNode(*n, *op_def, nullptr, &output_names));
+    for (const auto& o : output_names) {
+      CHECK_LE(o.second.second, src.ret.size())
+          << "Missing ret for output '" << o.first << "' in '" << src.ret[0]
+          << "' of " << name;
+      for (int i = o.second.first; i < o.second.second; ++i) {
+        ret_index[src.ret[i]] =
+            strings::StrCat(src.ret[0], ":", o.first, ":", i - o.second.first);
+      }
+    }
+  }
+
+  // Returns
+  for (const auto& r : fdef.signature().output_arg()) {
+    const auto iter = ret_index.find(r.name());
+    CHECK(iter != ret_index.end()) << "Return '" << r.name() << "' in " << name;
+    fdef.mutable_ret()->insert({r.name(), iter->second});
+  }
+  return fdef;
+}
+
+FunctionDef FunctionDefHelper::Define(gtl::ArraySlice<string> arg_def,
+                                      gtl::ArraySlice<string> ret_def,
+                                      gtl::ArraySlice<string> attr_def,
+                                      gtl::ArraySlice<Node> node_def) {
+  return Define("_", arg_def, ret_def, attr_def, node_def);
+}
+
+namespace gradient {
+
+typedef std::unordered_map<string, Creator> OpGradFactory;
+
+OpGradFactory* GetOpGradFactory() {
+  static OpGradFactory* factory = new OpGradFactory;
+  return factory;
+}
+
+bool RegisterOp(const string& op, Creator func) {
+  CHECK(GetOpGradFactory()->insert({op, func}).second)
+      << "Duplicated gradient for " << op;
+  return true;
+}
+
+Status GetOpGradientCreator(const string& op, Creator* creator) {
+  auto fac = GetOpGradFactory();
+  auto iter = fac->find(op);
+  if (iter == fac->end()) {
+    return errors::NotFound("No gradient defined for op: ", op);
+  }
+  *creator = iter->second;
+  return Status::OK();
+}
+
+}  // end namespace gradient
+
+}  // end namespace tensorflow

function.h

@@ -0,0 +1,625 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_FUNCTION_H_
+#define TENSORFLOW_FRAMEWORK_FUNCTION_H_
+
+#include <vector>
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/selective_registration.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+
+class CancellationManager;
+class GraphDef;
+class OpKernel;
+class ResourceMgr;
+class Rendezvous;
+class ScopedStepContainer;
+class StepStatsCollector;
+class Node;
+
+// FunctionDefHelper::Create is a convenient helper to construct a
+// FunctionDef proto.
+// E.g.,
+//   FunctionDef my_func = FunctionDefHelper::Create(
+//     "my_func_name",
+//     {"x:T", "y:T" /* one string per argument */},
+//     {"z:T" /* one string per return value */},
+//     {"T: {float, double}" /* one string per attribute  */},
+//     {
+//        {{"o"}, "Mul", {"x", "y"}, {{"T", "$T"}}}
+//        /* one entry per function node */
+//     },
+//     /* Mapping between function returns and function node outputs. */
+//     {{"z", "o:z"}});
+//
+// For the old Function::Node approach, use FunctionDefHelper::Define()
+// E.g.,
+//   FunctionDef my_func = FunctionDefHelper::Define(
+//     "my_func_name",
+//     {"x:T", "y:T" /* one string per argument */},
+//     {"z:T" /* one string per return value */},
+//     {"T: {float, double}" /* one string per attribute  */},
+//     {
+//        {{"z"}, "Mul", {"x", "y"}, {{"T", "$T"}}}
+//        /* one entry per function node */
+//     });
+class FunctionDefHelper {
+ public:
+  // AttrValueWrapper has copy constructors for the type T so that
+  // it's easy to construct a simple AttrValue proto.
+  //
+  // If T is a string type (const char*, string, or StringPiece), and
+  // it starts with "$", we construct a AttrValue of "placeholder".
+  //
+  // E.g.,
+  //   std::<string, AttrValueWrapper> x = {"T", "$T"}
+  // is a named attr value placeholder.
+  struct AttrValueWrapper {
+    AttrValue proto;
+
+    AttrValueWrapper() {}
+
+    template <typename T>
+    AttrValueWrapper(T val) {  // NOLINT(runtime/explicit)
+      SetAttrValue(val, &proto);
+    }
+
+   private:
+    void InitFromString(StringPiece val);
+  };
+
+  // Constructs an AttrValue.func given the "name" and "attrs".
+  static AttrValueWrapper FunctionRef(
+      const string& name,
+      gtl::ArraySlice<std::pair<string, AttrValueWrapper>> attrs);
+  static AttrValueWrapper FunctionRef(const string& name) {
+    return FunctionRef(name, {});
+  }
+
+  // Node is used to construct FunctionDef.Node using initialization
+  // lists. E.g.,
+  //  Node n = {{"z"}, "Mul", {"x", "y"}, {{"T", "$T"}}};  // z = x * y
+  struct Node {
+    // When constructing a NodeDef, the first entry in ret is used as
+    // the node name, the remaining values are ignored.
+    std::vector<string> ret;
+    string op;
+    std::vector<string> arg;
+    std::vector<std::pair<string, AttrValueWrapper>> attr;
+    std::vector<string> dep;
+
+    NodeDef ToNodeDef() const;
+  };
+
+  // The Create() function uses the new NodeDef field.  `ret_def`
+  // holds a mapping from the function output names from `out_def` to
+  // the node outputs from `node_def`.
+  static FunctionDef Create(const string& function_name,
+                            gtl::ArraySlice<string> in_def,
+                            gtl::ArraySlice<string> out_def,
+                            gtl::ArraySlice<string> attr_def,
+                            gtl::ArraySlice<Node> node_def,
+                            gtl::ArraySlice<std::pair<string, string>> ret_def);
+
+  // The two Define() functions use the old FunctionDef::Node field.
+  // TODO(josh11b): Get rid of these and transition to the one above.
+  static FunctionDef Define(const string& function_name,
+                            gtl::ArraySlice<string> arg_def,
+                            gtl::ArraySlice<string> ret_def,
+                            gtl::ArraySlice<string> attr_def,
+                            gtl::ArraySlice<Node> node_def);
+
+  // Defines an anonymous function. I.e., its name is not relevant.
+  static FunctionDef Define(gtl::ArraySlice<string> arg_def,
+                            gtl::ArraySlice<string> ret_def,
+                            gtl::ArraySlice<string> attr_def,
+                            gtl::ArraySlice<Node> node_def);
+
+  // Helpers to construct a constant scalar.
+  template <typename T>
+  static Node Const(const string& name, const T& val) {
+    Node n = {{name}, "Const"};
+    const DataType dtype = DataTypeToEnum<T>::value;
+    n.attr.push_back({"dtype", dtype});
+    Tensor t(dtype, TensorShape({}));
+    t.scalar<T>()() = val;
+    n.attr.push_back({"value", t});
+    return n;
+  }
+
+  template <typename T>
+  static Node Const(const string& name, gtl::ArraySlice<T> vals) {
+    Node n = {{name}, "Const"};
+    const DataType dtype = DataTypeToEnum<T>::value;
+    n.attr.push_back({"dtype", dtype});
+    int64 num = vals.size();
+    Tensor t(dtype, TensorShape({num}));
+    for (size_t i = 0; i < vals.size(); ++i) {
+      t.flat<T>()(i) = vals[i];
+    }
+    n.attr.push_back({"value", t});
+    return n;
+  }
+};
+
+template <>
+inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper(const char* val) {
+  InitFromString(val);
+}
+
+template <>
+inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper(
+    const string& val) {
+  InitFromString(val);
+}
+
+template <>
+inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper(StringPiece val) {
+  InitFromString(val);
+}
+
+// Instantiate a function.
+//
+// "fdef" encodes a TF function with some attrs in fdef.signature.attr
+// containing placeholders.  InstantiateFunction binds these
+// placeholders and produces an instantiated function encoded in
+// "result.gdef". The value to substitute a placeholder is given by
+// "attr_values", which is a map from a placeholder name to an attr
+// value.
+//
+// InstantiateFunction calls "get_function" to find signatures of other
+// functions and primitive ops.
+
+// GetFunctionSignature(func name, opdef) returns OK if the func name is found
+// and opdef is filled with a pointer to the corresponding signature
+// (a OpDef proto). Otherwise, returns an error.
+typedef std::function<Status(const string&, const OpDef**)>
+    GetFunctionSignature;
+
+struct InstantiationResult {
+  DataTypeVector arg_types;
+  DataTypeVector ret_types;
+  std::vector<NodeDef> nodes;
+};
+Status InstantiateFunction(const FunctionDef& fdef, AttrSlice attr_values,
+                           GetFunctionSignature get_function,
+                           InstantiationResult* result);
+
+// Returns a debug string for a function definition.
+//
+// The returned text is multiple-line. It is intended to be
+// human-readable rather than being friendly to parsers. It is _NOT_
+// intended to be the canonical string representation of "func_def".
+// Particularly, it may not include all information presented in
+// "func_def" (e.g., comments, description of the function arguments,
+// etc.)
+string DebugString(const FunctionDef& func_def);
+string DebugString(const GraphDef& instantiated_func_def);
+string DebugString(gtl::ArraySlice<NodeDef> instantiated_func_nodes);
+
+// Returns a debug string for a top level graph (the main program and
+// its supporting functions defined in its library).
+string DebugStringWhole(const GraphDef& gdef);
+
+// Returns true if f1 == f2. Compares all fields, including descriptions. Order
+// of NodeDefs doesn't matter.
+bool FunctionDefsEqual(const FunctionDef& f1, const FunctionDef& f2);
+
+// Return a hash of `fdef` that is consistent with FunctionDefsEqual method.
+// In other words, if two fdefs compare equal, their hash values will be the
+// same.
+uint64 FunctionDefHash(const FunctionDef& fdef);
+
+// Returns a canonicalized string for the instantiation of the
+// function of the given "name" and attributes "attrs".
+//
+// The returned string is guaranteed to be stable within one address
+// space. But it may be change as the implementation
+// evolves. Therefore, it should not be persisted or compared across
+// address spaces.
+string Canonicalize(const string& funcname, AttrSlice attrs);
+
+class CallFrameInterface {
+ public:
+  virtual ~CallFrameInterface() {}
+
+  virtual size_t num_args() const = 0;
+  virtual size_t num_retvals() const = 0;
+
+  virtual Status GetArg(int index, Tensor* val) const = 0;
+  virtual Status SetRetval(int index, const Tensor& val) = 0;
+};
+
+// Represents a function call frame. I.e., the data structure used to
+// pass arguments to a function and retrieve its results.
+//
+// Runtime must arrange accesses to one FunctionCallFrame s.t.
+//   1. SetArgs() happens before any GetArg();
+//   2. GetRetvals happens after all SetRetval();
+class FunctionCallFrame : public CallFrameInterface {
+ public:
+  FunctionCallFrame(DataTypeSlice arg_types, DataTypeSlice ret_types);
+  ~FunctionCallFrame();
+
+  // Caller methods.
+  Status SetArgs(gtl::ArraySlice<Tensor> args);
+  Status GetRetvals(std::vector<Tensor>* rets) const;
+  Status ConsumeRetvals(std::vector<Tensor>* rets);
+
+  size_t num_args() const override { return arg_types_.size(); }
+  size_t num_retvals() const override { return ret_types_.size(); }
+
+  // Callee methods.
+  Status GetArg(int index, Tensor* val) const override;
+  Status SetRetval(int index, const Tensor& val) override;
+
+ private:
+  DataTypeVector arg_types_;
+  DataTypeVector ret_types_;
+  gtl::InlinedVector<Tensor, 4> args_;
+  struct Retval {
+    bool has_val = false;
+    Tensor val;
+  };
+  gtl::InlinedVector<Retval, 4> rets_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(FunctionCallFrame);
+};
+
+// Helper to maintain a map between function names in a given
+// FunctionDefLibrary and function definitions.
+class FunctionLibraryDefinition : public OpRegistryInterface {
+ public:
+  explicit FunctionLibraryDefinition(const FunctionLibraryDefinition& lib_def);
+  FunctionLibraryDefinition(const OpRegistryInterface* default_registry,
+                            const FunctionDefLibrary& lib_def);
+  ~FunctionLibraryDefinition() override;
+
+  FunctionLibraryDefinition& operator=(const FunctionLibraryDefinition&) =
+      delete;
+
+  // Returns nullptr if "func" is not defined in "lib_def". Otherwise,
+  // returns its definition proto.
+  const FunctionDef* Find(const string& func) const;
+
+  // Adds function definition 'fdef' to this function library.
+  // Returns status 'ok' on success, or error otherwise. This is a no-op if
+  // 'fdef' already exists in this function library.
+  // If 'fdef' is successfully added to the library, it will be accessible
+  // from 'LookUp' and included in the proto returned by 'ToProto'.
+  // This operation is atomic.
+  Status AddFunctionDef(const FunctionDef& fdef);
+
+  // Adds gradient definition 'grad' to this function library.
+  // This is a no-op if 'grad' already exists in this function library.
+  // If 'grad' is successfully added, it will be accessible via 'FindGradient'
+  // and included in the proto returned by 'ToProto'.
+  // This operation is atomic.
+  Status AddGradientDef(const GradientDef& grad);
+
+  // Adds the functions and gradients in 'other' to this function library.
+  // Duplicate functions and gradients are ignored.
+  // This operation is atomic.
+  Status AddLibrary(const FunctionLibraryDefinition& other);
+
+  // Adds the functions and gradients in 'lib_def' to this function library.
+  // Duplicate functions and gradients are ignored.
+  // This operation is atomic.
+  Status AddLibrary(const FunctionDefLibrary& lib_def);
+
+  // If the gradient function for 'func' is specified explicitly in
+  // the library, returns the gradient function name.  Otherwise,
+  // returns an empty string.
+  string FindGradient(const string& func) const;
+
+  // OpRegistryInterface method. Useful for constructing a Graph.
+  //
+  // If "op" is defined in the library, returns its signature.
+  // Otherwise, assume "op" is a primitive op and returns its op
+  // signature and shape inference function.
+  Status LookUp(const string& op_type_name,
+                const OpRegistrationData** op_reg_data) const override;
+
+  static constexpr const char* const kGradientOp = "SymbolicGradient";
+  static constexpr const char* const kFuncAttr = "f";
+
+  // Given a node def 'ndef', inspects attributes of the callee
+  // function to derive the attribute 'value' for 'attr'. Returns OK
+  // iff the attribute is given by the function's definition.
+  // TODO(irving): Remove; keep only the const Node& version.
+  template <typename T>
+  Status GetAttr(const NodeDef& ndef, const string& attr, T* value) const;
+
+  // Given a node, inspects attributes of the callee function to derive the
+  // attribute 'value' for 'attr'. Returns OK iff the attribute is given by the
+  // function's definition.
+  template <typename T>
+  Status GetAttr(const Node& node, const string& attr, T* value) const;
+
+  // Returns a proto representation of the state of this function library.
+  FunctionDefLibrary ToProto() const;
+
+  const OpRegistryInterface* default_registry() const {
+    return default_registry_;
+  }
+
+ private:
+  // Shape inference for functions is handled separately by ShapeRefiner.
+
+  struct FunctionDefAndOpRegistration {
+    FunctionDefAndOpRegistration(const FunctionDef& fdef_in);
+
+    FunctionDef fdef;
+    OpRegistrationData op_registration_data;
+  };
+
+  // Same as AddFunctionDef/AddGradientDef except these methods set
+  // `added` to true if the `fdef`/`grad` were actually added to this.
+  Status AddFunctionDefHelper(const FunctionDef& fdef, bool* added);
+  Status AddGradientDefHelper(const GradientDef& grad, bool* added);
+
+  const OpRegistryInterface* const default_registry_;
+  gtl::FlatMap<string, std::unique_ptr<FunctionDefAndOpRegistration>>
+      function_defs_;
+  gtl::FlatMap<string, string> func_grad_;
+
+  // Helper function for GetAttr. Returns the FunctionDef* to get the
+  // attr from.
+  const FunctionDef* GetAttrImpl(const NodeDef& ndef) const;
+
+  // Remove function `func` from the library. `func` must be in the library.
+  void RemoveFunction(const string& func);
+
+  // Remove gradient of function `func` from the library. `func` must have
+  // a gradient.
+  void RemoveGradient(const string& func);
+
+  // Remove all functions in `funcs` and all gradients of
+  // functions in `funcs_with_grads` from this library.
+  void Remove(const std::vector<string>& funcs,
+              const std::vector<string>& funcs_with_grads);
+};
+
+// Forward declare. Defined in common_runtime/function.h
+struct FunctionBody;
+
+// Forward declare. Defined in common_runtime/device.h
+class Device;
+
+class FunctionLibraryRuntime {
+ public:
+  virtual ~FunctionLibraryRuntime() {}
+
+  // Instantiate a function with the given "attrs".
+  //
+  // Returns OK and fills in "handle" if the instantiation succeeds.
+  // Otherwise returns an error and "handle" is undefined.
+  typedef uint64 Handle;
+  virtual Status Instantiate(const string& function_name, AttrSlice attrs,
+                             Handle* handle) = 0;
+
+  // Releases state associated with the handle.
+  virtual Status ReleaseHandle(Handle handle) = 0;
+
+  // Returns the function body for the instantiated function given its
+  // handle 'h'. Returns nullptr if "h" is not found.
+  //
+  // *this keeps the ownership of the returned object, which remains alive
+  // as long as *this.
+  virtual const FunctionBody* GetFunctionBody(Handle h) = 0;
+
+  // Asynchronously invokes the instantiated function identified by
+  // "handle".
+  //
+  // If function execution succeeds, "done" is called with OK and
+  // "*rets" is filled with the function's return values. Otheriwse,
+  // "done" is called with an error status.
+  //
+  // Does not take ownership of "rets".
+  // In the cross-process scenario, runner isn't used for making the Async
+  // RPC calls.
+  struct Options {
+    // The id of the step that is calling this function.
+    int64 step_id = 0;
+    Rendezvous* rendezvous = nullptr;
+    CancellationManager* cancellation_manager = nullptr;
+    ScopedStepContainer* step_container = nullptr;
+    StepStatsCollector* stats_collector = nullptr;
+
+    std::function<void(std::function<void()>)>* runner = nullptr;
+
+    // Parameters for remote function execution.
+    bool remote_execution = false;
+    string source_device = "";  // Fully specified device name.
+
+    // Allocator attributes specifying where the args are / rets should be put.
+    // These should either be {} or match the length of args / retvals. If {},
+    // the default allocator attributes will be assumed for all args / retvals.
+    std::vector<AllocatorAttributes> args_alloc_attrs;
+    std::vector<AllocatorAttributes> rets_alloc_attrs;
+
+    // If true, we create a new IntraProcessRendezvous, else use the existing
+    // one.
+    bool create_rendezvous = false;
+  };
+  typedef std::function<void(const Status&)> DoneCallback;
+  virtual void Run(const Options& opts, Handle handle,
+                   gtl::ArraySlice<Tensor> args, std::vector<Tensor>* rets,
+                   DoneCallback done) = 0;
+  virtual void Run(const Options& opts, Handle handle,
+                   CallFrameInterface* call_frame, DoneCallback done) = 0;
+
+  // Creates a "kernel" for the given node def "ndef".
+  //
+  // If succeeds, returns OK and the caller takes the ownership of the
+  // returned "*kernel". Otherwise, returns an error.
+  virtual Status CreateKernel(const NodeDef& ndef, OpKernel** kernel) = 0;
+
+  // Returns true iff 'function' is stateful.
+  virtual bool IsStateful(const string& function_name) = 0;
+
+  // Returns the device on which the function executes.
+  virtual Device* device() = 0;
+
+  // Returns the function library definition that backs this runtime.
+  virtual const FunctionLibraryDefinition* GetFunctionLibraryDefinition()
+      const = 0;
+
+  // Returns the environment on which the function executes.
+  virtual Env* env() = 0;
+
+  // Returns a debug string showing the definition of the function of
+  // 'handle'.
+  virtual string DebugString(Handle handle) = 0;
+
+  // Returns the graph version number.
+  virtual int graph_def_version() = 0;
+
+  typedef uint64 LocalHandle;
+};
+
+const FunctionLibraryRuntime::Handle kInvalidHandle = -1;
+const FunctionLibraryRuntime::LocalHandle kInvalidLocalHandle = -1;
+typedef std::function<Status(FunctionLibraryRuntime*, const NodeDef&,
+                             std::unique_ptr<OpKernel>*)>
+    CustomKernelCreator;
+
+// Used to instantiate and run functions in a distributed system.
+class DistributedFunctionLibraryRuntime {
+ public:
+  virtual ~DistributedFunctionLibraryRuntime() {}
+
+  // The _target attr in attrs determines where the function is instantiated.
+  virtual Status Instantiate(const string& function_name,
+                             const FunctionLibraryDefinition& lib_def,
+                             AttrSlice attrs,
+                             FunctionLibraryRuntime::LocalHandle* handle) = 0;
+
+  // opts.runner isn't used for execution.
+  virtual void Run(const FunctionLibraryRuntime::Options& opts,
+                   FunctionLibraryRuntime::LocalHandle handle,
+                   gtl::ArraySlice<Tensor> args, std::vector<Tensor>* rets,
+                   FunctionLibraryRuntime::DoneCallback done) = 0;
+};
+
+// Extracts the actual type from "attr_values" based on its definition
+// "arg_def".
+//
+// If "arg_def" is a N*T type, *is_type_list is set to false, and
+// *dtypes is set to be a vector of size N and each element is T.
+//
+// If "arg_def" is a list(type), *is_type_list is set to true, and
+// *dtypes is set to be a vector of types specified in attrs for
+// arg_def.
+//
+// Otherwise (arg_def is a simple type T), *is_type_list is set to
+// false, and *dtypes is set to a single element vector, whose only
+// element is T.
+Status ArgNumType(AttrSlice attrs, const OpDef::ArgDef& arg_def,
+                  bool* is_type_list, DataTypeVector* dtypes);
+
+// To register a gradient function for a builtin op, one should use
+//   REGISTER_OP_GRADIENT(<op_name>, <c++ grad factory>);
+//
+// Typically, the c++ grad factory is a plan function that can be
+// converted into ::tensorflow::gradient::Creator, which is
+//   std::function<Status(const AttrSlice&, FunctionDef*)>.
+//
+// A ::tensorflow::gradient::Creator should populate in FunctionDef* with a
+// definition of a brain function which compute the gradient for the
+// <op_name> when the <op_name> is instantiated with the given attrs.
+//
+// E.g.,
+//
+// Status MatMulGrad(const AttrSlice& attrs, FunctionDef* g) {
+//   bool transpose_a;
+//   TF_RETURN_IF_ERROR(attrs.Get("transpose_a", &transpose_a));
+//   bool transpose_b;
+//   TF_RETURN_IF_ERROR(attrs.Get("transpose_b", &transpose_b));
+//   DataType dtype;
+//   TF_RETURN_IF_ERROR(attrs.Get("dtype", &dtype));
+//   if (!transpose_a && !transpose_b) {
+//     *g = FunctionDefHelper::Define(
+//       "MatMulGrad",
+//       {"x:T ", "y:T", "dz:T"},    // Inputs to this function
+//       {"dx:T", "dy:T"},           // Outputs from this function
+//       {"T: {float, double}"},     // Attributes needed by this function
+//       {
+//         {{"x_t"}, "Transpose", {"x"}, {{"T", "$T"}}},
+//         {{"y_t"}, "Transpose", {"y"}, {{"T", "$T"}}},
+//         {{"dx"}, "MatMul", {"dz", "y_t"}, {{"T", "$T"}}},
+//         {{"dy"}, "MatMul", {"x_", "dz"}, {{"T", "$T"}}},
+//       });
+//   } else {
+//     ... ...
+//   }
+//   return Status::OK();
+// }
+//
+// NOTE: $T is substituted with the type variable "T" when the
+// gradient function MatMul is instantiated.
+//
+// TODO(zhifengc): Better documentation somewhere.
+
+// Macros to define a gradient function factory for a primitive
+// operation.
+#define REGISTER_OP_GRADIENT(name, fn) \
+  REGISTER_OP_GRADIENT_UNIQ_HELPER(__COUNTER__, name, fn)
+
+#define REGISTER_OP_NO_GRADIENT(name) \
+  REGISTER_OP_GRADIENT_UNIQ_HELPER(__COUNTER__, name, nullptr)
+
+#define REGISTER_OP_GRADIENT_UNIQ_HELPER(ctr, name, fn) \
+  REGISTER_OP_GRADIENT_UNIQ(ctr, name, fn)
+
+#define REGISTER_OP_GRADIENT_UNIQ(ctr, name, fn)                 \
+  static bool unused_grad_##ctr = SHOULD_REGISTER_OP_GRADIENT && \
+                                  ::tensorflow::gradient::RegisterOp(name, fn)
+
+namespace gradient {
+// Register a gradient creator for the "op".
+typedef std::function<Status(const AttrSlice& attrs, FunctionDef*)> Creator;
+bool RegisterOp(const string& op, Creator func);
+
+// Returns OK the gradient creator for the "op" is found (may be
+// nullptr if REGISTER_OP_NO_GRADIENT is used.
+Status GetOpGradientCreator(const string& op, Creator* creator);
+};
+
+// Declare explicit instantiations of GetAttr
+#define GET_ATTR(T)                                          \
+  extern template Status FunctionLibraryDefinition::GetAttr( \
+      const Node&, const string&, T*) const;                 \
+  extern template Status FunctionLibraryDefinition::GetAttr( \
+      const NodeDef&, const string&, T*) const;
+GET_ATTR(string)
+GET_ATTR(bool)
+#undef GET_ATTR
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_FUNCTION_H_

function.proto

@@ -0,0 +1,101 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "FunctionProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/attr_value.proto";
+import "tensorflow/core/framework/node_def.proto";
+import "tensorflow/core/framework/op_def.proto";
+
+// A library is a set of named functions.
+message FunctionDefLibrary {
+  repeated FunctionDef function = 1;
+  repeated GradientDef gradient = 2;
+}
+
+// A function can be instantiated when the runtime can bind every attr
+// with a value. When a GraphDef has a call to a function, it must
+// have binding for every attr defined in the signature.
+//
+// TODO(zhifengc):
+//   * device spec, etc.
+message FunctionDef {
+  // The definition of the function's name, arguments, return values,
+  // attrs etc.
+  OpDef signature = 1;
+
+  // Attributes specific to this function definition.
+  map<string, AttrValue> attr = 5;
+
+  // NOTE: field id 2 deleted on Jan 11, 2016, GraphDef version 21.
+
+  // In both of the following fields, there is the need to specify an
+  // output that is used as either the input to another node (in
+  // `node_def`) or as a return value of the function (in `ret`).
+  // Unlike the NodeDefs in GraphDef, we need to be able to specify a
+  // list in some cases (instead of just single outputs).  Also, we
+  // need to be able to deal with lists of unknown length (so the
+  // output index may not be known at function definition time).  So
+  // we use the following format instead:
+  // * "fun_in" where "fun_in" is the name of a function input arg in
+  //   the `signature` field above.  This represents that input, whether
+  //   it is a single tensor or a list.
+  // * "fun_in:0" gives the first element of a function input arg (a
+  //   non-list input is considered a list of length 1 for these
+  //   purposes).
+  // * "node:out" where "node" is the name of a node in `node_def` and
+  //   "out" is the name one of its op's output arguments (the name
+  //   comes from the OpDef of the node's op). This represents that
+  //   node's output, whether it is a single tensor or a list.
+  //   Note: We enforce that an op's output arguments are never
+  //   renamed in the backwards-compatibility test.
+  // * "node:out:0" gives the first element of a node output arg (a
+  //   non-list output is considered a list of length 1 for these
+  //   purposes).
+  //
+  // NOT CURRENTLY SUPPORTED (but may be in the future):
+  // * "node:out:-1" gives last element in a node output list
+  // * "node:out:1:" gives a list with all but the first element in a
+  //   node output list
+  // * "node:out::-1" gives a list with all but the last element in a
+  //   node output list
+
+  // The body of the function.  Unlike the NodeDefs in a GraphDef, attrs
+  // may have values of type `placeholder` and the `input` field uses
+  // the "output" format above.
+
+  // By convention, "op" in node_def is resolved by consulting with a
+  // user-defined library first. If not resolved, "func" is assumed to
+  // be a builtin op.
+  repeated NodeDef node_def = 3;
+
+  // A mapping from the output arg names from `signature` to the
+  // outputs from `node_def` that should be returned by the function.
+  map<string, string> ret = 4;
+}
+
+// GradientDef defines the gradient function of a function defined in
+// a function library.
+//
+// A gradient function g (specified by gradient_func) for a function f
+// (specified by function_name) must follow the following:
+//
+// The function 'f' must be a numerical function which takes N inputs
+// and produces M outputs. Its gradient function 'g', which is a
+// function taking N + M inputs and produces N outputs.
+//
+// I.e. if we have
+//    (y1, y2, ..., y_M) = f(x1, x2, ..., x_N),
+// then, g is
+//    (dL/dx1, dL/dx2, ..., dL/dx_N) = g(x1, x2, ..., x_N,
+//                                      dL/dy1, dL/dy2, ..., dL/dy_M),
+// where L is a scalar-value function of (x1, x2, ..., xN) (e.g., the
+// loss function). dL/dx_i is the partial derivative of L with respect
+// to x_i.
+message GradientDef {
+  string function_name = 1;  // The function name.
+  string gradient_func = 2;  // The gradient function's name.
+}

function_test.cc

@@ -0,0 +1,1339 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/function.h"
+#include <vector>
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+// A helper class to make AttrSlice from initializer lists
+class Attrs {
+ public:
+  Attrs(const std::initializer_list<  // NOLINT(runtime/explicit)
+        std::pair<string, FunctionDefHelper::AttrValueWrapper>>
+            attrs) {
+    for (const auto& aval : attrs) {
+      map_.insert({aval.first, aval.second.proto});
+    }
+  }
+
+  operator AttrSlice() { return AttrSlice(&map_); }  // NOLINT(runtime/explicit)
+
+ private:
+  AttrValueMap map_;
+};
+
+typedef FunctionDefHelper FDH;
+
+Status GetOpSig(const string& op, const OpDef** sig) {
+  return OpRegistry::Global()->LookUpOpDef(op, sig);
+}
+
+REGISTER_OP("One")
+    .Output("y: T")
+    .Attr("T: {float, double, int32, int64}")
+    .Doc(R"doc(
+Returns a tensor with a single element (1) of type T.
+
+y: A scalar in type T.
+
+)doc");
+
+TEST(TFunc, SquarePlusOne) {
+  auto fdef = FDH::Create(
+      // Name
+      "SquarePlusOne",
+      // Inputs
+      {"x: T"},
+      // Outputs
+      {"y: T"},
+      // Attrs
+      {"T: {float, double, int32, int64}"},
+      // Nodes
+      {// a = Square<T>(x)
+       {{"a"}, "Square", {"x"}, {{"T", "$T"}}},
+       // o = One<T>()
+       // NOTE: We can also have a Cast<Tin, Tout>(x) instead.
+       {{"o"}, "One", {}, {{"T", "$T"}}},
+       // y = Add<T>(a, o)
+       {{"y"}, "Add", {"a:y", "o:y"}, {{"T", "$T"}}}},
+      // Returns
+      {{"y", "y:z:0"}});
+
+  const char* e = R"P(
+SquarePlusOne[T:{float, double, int32, int64}](x:T) -> (y:T) {
+  a = Square[T=$T](x)
+  o = One[T=$T]()
+  y = Add[T=$T](a:y, o:y)
+  return y = y:z:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  // Instantiate one with T=float
+  InstantiationResult result;
+  TF_ASSERT_OK(
+      InstantiateFunction(fdef, Attrs({{"T", DT_FLOAT}}), GetOpSig, &result));
+  const char* e2 = R"P(
+(x:float) -> (y:float) {
+  a = Square[T=float](x)
+  o = One[T=float]()
+  y = Add[T=float](a, o)
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+TEST(TFunc, ControlDep) {
+  auto fdef = FDH::Create(
+      // Name
+      "ControlDep",
+      // Inputs
+      {"x: int32"},
+      // Outputs
+      {"y: int32"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = Identity<int32>(x)
+       {{"a"}, "Identity", {"x"}, {{"T", DT_INT32}}},
+       // o = NoOp(^a)
+       {{"o"}, "NoOp", {"^a"}, {}},
+       // y = Identity<int32>(a, ^o)
+       {{"y"}, "Identity", {"a:output:0", "^o"}, {{"T", DT_INT32}}}},
+      // Returns
+      {{"y", "y:output:0"}});
+
+  const char* e = R"P(
+ControlDep(x:int32) -> (y:int32) {
+  a = Identity[T=int32](x)
+  o = NoOp() @ a
+  y = Identity[T=int32](a:output:0) @ o
+  return y = y:output:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  // Instantiate one with T=float
+  InstantiationResult result;
+  TF_ASSERT_OK(
+      InstantiateFunction(fdef, Attrs({{"T", DT_FLOAT}}), GetOpSig, &result));
+  const char* e2 = R"P(
+(x:int32) -> (y:int32) {
+  a = Identity[T=int32](x)
+  o = NoOp() @ a
+  y = Identity[T=int32](a) @ o
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_INT32}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_INT32}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+REGISTER_OP("HasDefaultType")
+    .Output("out: T")
+    .Attr("T: {float, double, int32, int64} = DT_FLOAT");
+
+// This verifies that a function using an op before a type attr (with
+// a default) is added, still works.  This is important for backwards
+// compatibility.
+TEST(TFunc, MissingTypeAttr) {
+  auto fdef = FDH::Create(
+      // Name
+      "BackCompat",
+      // Args
+      {},
+      // Return values
+      {"y: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// y = HasDefaultType(x), T missing, defaults to float
+       {{"a"}, "HasDefaultType", {}, {}}},
+      // Returns
+      {{"y", "a:out:0"}});
+
+  const char* e = R"P(
+BackCompat() -> (y:float) {
+  a = HasDefaultType()
+  return y = a:out:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result));
+  // Should get T=float from Op's default.
+  const char* e2 = R"P(
+() -> (a:float) {
+  a = HasDefaultType[T=float]()
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector());
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+TEST(TFunc, NTimesT) {
+  auto fdef = FDH::Create(
+      // Name
+      "NTimesT",
+      // Inputs
+      {"x: float", "y: float"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = AddN<N=2>(x, y)
+       {{"a"}, "AddN", {"x", "y"}, {{"T", DT_FLOAT}, {"N", 2}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  const char* e = R"P(
+NTimesT(x:float, y:float) -> (z:float) {
+  a = AddN[N=2, T=float](x, y)
+  return z = a:sum:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result));
+  const char* e2 = R"P(
+(x:float, y:float) -> (a:float) {
+  a = AddN[N=2, T=float](x, y)
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT, DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+// NOTE: This is the simplest Map op. It takes a f:T->U.
+REGISTER_OP("Map")
+    .Input("x: N * T")
+    .Output("y: N * U")
+    .Attr("T: type")
+    .Attr("U: type")
+    .Attr("N: int >= 1")
+    // .Attr("func: func_name_with_attr")
+    .Doc(R"doc(
+Applies the 'func' on every input. I.e.,
+
+y[i] = func<...>(x[i])
+
+x: N tensors, each of type T;
+y: N tensors, each of type U;
+
+)doc");
+
+TEST(TFunc, AddSquared) {
+  auto fdef = FDH::Create(
+      // Name
+      "AddSquared",
+      // Args
+      {"x: N*T"},
+      // Return values
+      {"y: T"},
+      // Attrs
+      {"N:int", "T:{float, double, int32, int64}"},
+      // Nodes
+      {// a = Map<func=Square<$T>,T=$T,U=$T,N=$N>(x)
+       {{"a"},
+        "Map",
+        {"x"},
+        {{"func", FDH::FunctionRef("Square", {{"T", "$T"}})},
+         {"T", "$T"},
+         {"U", "$T"},
+         {"N", "$N"}}},
+       // y = AddN<N=$N,T=$T>(a)
+       {{"y"}, "AddN", {"a:y"}, {{"N", "$N"}, {"T", "$T"}}}},
+      {{"y", "y:sum"}});
+
+  const char* e = R"P(
+AddSquared[N:int, T:{float, double, int32, int64}](x:N*T) -> (y:T) {
+  a = Map[N=$N, T=$T, U=$T, func=Square[T=$T]](x)
+  y = AddN[N=$N, T=$T](a:y)
+  return y = y:sum
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  // Instantiate one with T=float
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, Attrs({{"N", 3}, {"T", DT_FLOAT}}),
+                                   GetOpSig, &result));
+  const char* e2 = R"P(
+(x_0:float, x_1:float, x_2:float) -> (y:float) {
+  a = Map[N=3, T=float, U=float, func=Square[T=float]](x_0, x_1, x_2)
+  y = AddN[N=3, T=float](a, a:1, a:2)
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT, DT_FLOAT, DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+TEST(TFunc, ControlDeps) {
+  auto fdef = FDH::Define(
+      // Name
+      "ControlDeps",
+      // Args
+      {"x: float"},
+      // Return values
+      {},
+      // Attrs
+      {},
+      // Nodes
+      {
+          {{"a"}, "One", {}, {{"T", DT_FLOAT}}, {"x"}},
+          {{"u"}, "NoOp", {}, {}, {"a"}},
+          {{"b"}, "One", {}, {{"T", DT_FLOAT}}, {"u"}},
+          {{"v"}, "NoOp", {}, {}, {"b"}},
+          {{"c"}, "One", {}, {{"T", DT_FLOAT}}, {"a", "v"}},
+      });
+  const char* e = R"P(
+ControlDeps(x:float) -> () {
+  a = One[T=float]() @ x
+  u = NoOp() @ a
+  b = One[T=float]() @ u
+  v = NoOp() @ b
+  c = One[T=float]() @ a, v
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result));
+  const char* e2 = R"P(
+(x:float) -> () {
+  a = One[T=float]() @ x
+  u = NoOp() @ a
+  b = One[T=float]() @ u
+  v = NoOp() @ b
+  c = One[T=float]() @ a, v
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+TEST(TFunc, XTimesTwo) {
+  auto expect = R"P(
+XTimesTwo[T:{float, double, int32, int64}](x:T) -> (y:T) {
+  two = Const[dtype=int64, value=Tensor<type: int64 shape: [] values: 2>]()
+  scale = Cast[DstT=$T, SrcT=int64](two:output:0)
+  y = Mul[T=$T](x, scale:y:0)
+  return y = y:z:0
+}
+)P";
+  EXPECT_EQ(expect, DebugString(test::function::XTimesTwo()));
+}
+
+TEST(TFunc, WXPlusB) {
+  auto expect = R"P(
+WXPlusB[T:{float, double}](w:T, x:T, b:T) -> (y:T) {
+  mm = MatMul[T=$T, _kernel="eigen", transpose_a=false, transpose_b=false](w, x)
+  y = Add[T=$T](mm:product:0, b)
+  return y = y:z:0
+}
+)P";
+  EXPECT_EQ(expect, DebugString(test::function::WXPlusB()));
+}
+
+TEST(TFunc, Body_TypeList) {
+  const Tensor kZero = test::AsScalar<int32>(0);
+  auto fdef = FDH::Create(
+      // Name
+      "Test",
+      // Args
+      {"i:float"},
+      // Return values
+      {"o:float"},
+      // Attrs
+      {},
+      // Nodes
+      {{{"zero"}, "Const", {}, {{"value", kZero}, {"dtype", DT_INT32}}},
+       {{"s"},
+        "Split",
+        {"zero:output:0", "i"},
+        {{"num_split", 4}, {"T", DT_FLOAT}}},
+       {{"l"}, "Mul", {"s:output:0", "s:output:1"}, {{"T", DT_FLOAT}}},
+       {{"r"}, "Mul", {"s:output:2", "s:output:3"}, {{"T", DT_FLOAT}}},
+       {{"x"},
+        "_ListToArray",
+        {"l:z", "r:z"},
+        {{"N", 2},
+         {"T", DT_FLOAT},
+         {"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}}}},
+       {{"o"}, "AddN", {"x:output"}, {{"N", 2}, {"T", DT_FLOAT}}}},
+      {{"o", "o:sum:0"}});
+
+  const char* e = R"P(
+Test(i:float) -> (o:float) {
+  zero = Const[dtype=int32, value=Tensor<type: int32 shape: [] values: 0>]()
+  s = Split[T=float, num_split=4](zero:output:0, i)
+  l = Mul[T=float](s:output:0, s:output:1)
+  r = Mul[T=float](s:output:2, s:output:3)
+  x = _ListToArray[N=2, T=float, Tin={float, float}](l:z, r:z)
+  o = AddN[N=2, T=float](x:output)
+  return o = o:sum:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result));
+  const char* e2 = R"P(
+(i:float) -> (o:float) {
+  zero = Const[dtype=int32, value=Tensor<type: int32 shape: [] values: 0>]()
+  s = Split[T=float, num_split=4](zero, i)
+  l = Mul[T=float](s, s:1)
+  r = Mul[T=float](s:2, s:3)
+  x = _ListToArray[N=2, T=float, Tin={float, float}](l, r)
+  o = AddN[N=2, T=float](x, x:1)
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+REGISTER_OP("Cond")
+    .Input("input: Tin")
+    .Output("output: out_types")
+    .Attr("Tin: list(type)")
+    .Attr("out_types: list(type)")
+    .Attr("cond: func")
+    .Attr("then_branch: func")
+    .Attr("else_branch: func")
+    .Doc(R"doc(
+output = Cond(input) ? then_branch(input) : else_branch(input)
+
+cond: A function takes 'input' and returns a scalar.
+then_branch: A function takes 'input' and returns 'output'.
+else_branch: A function takes 'input' and returns 'output'.
+)doc");
+
+TEST(TFunc, Body_Array_List_Converter) {
+  auto fdef = FDH::Define(
+      // Name
+      "MySelect",
+      // Args
+      {"x:float"},
+      // Return values
+      {"z:float"},
+      // Attrs
+      {},
+      // Nodes
+      {
+          {{"y"},
+           "Cond",
+           {"x"},
+           {{"Tin", DataTypeSlice{DT_FLOAT}},
+            {"out_types", DataTypeSlice{DT_FLOAT}},
+            {"cond", FDH::FunctionRef("MyCond")},
+            {"then_branch", FDH::FunctionRef("MyThen")},
+            {"else_branch", FDH::FunctionRef("MyElse")}}},
+          {{"z"},
+           "Cond",
+           {"y", "y"},
+           {{"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+            {"out_types", DataTypeSlice{DT_FLOAT}},
+            {"cond", FDH::FunctionRef("MyCond2")},
+            {"then_branch", FDH::FunctionRef("MyThen2")},
+            {"else_branch", FDH::FunctionRef("MyElse2")}}},
+      });
+
+  const char* e = R"P(
+MySelect(x:float) -> (z:float) {
+  y = Cond[Tin={float}, cond=MyCond, else_branch=MyElse, out_types={float}, then_branch=MyThen](x)
+  z = Cond[Tin={float, float}, cond=MyCond2, else_branch=MyElse2, out_types={float}, then_branch=MyThen2](y:output:0, y:output:0)
+  return z = z:output:0
+}
+)P";
+  EXPECT_EQ(DebugString(fdef), e);
+
+  InstantiationResult result;
+  TF_ASSERT_OK(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result));
+  const char* e2 = R"P(
+(x:float) -> (z:float) {
+  y = Cond[Tin={float}, cond=MyCond, else_branch=MyElse, out_types={float}, then_branch=MyThen](x)
+  z = Cond[Tin={float, float}, cond=MyCond2, else_branch=MyElse2, out_types={float}, then_branch=MyThen2](y, y)
+}
+)P";
+  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
+  EXPECT_EQ(DebugString(result.nodes), e2);
+}
+
+static void HasError(const Status& s, const string& substr) {
+  EXPECT_TRUE(StringPiece(s.ToString()).contains(substr))
+      << ">>" << s << "<<, expected substring >>" << substr << "<<";
+}
+
+TEST(InstantiateErrors, Not_Sufficient_Attrs) {
+  auto fdef =
+      FDH::Define("nop", {}, {}, {"T:{float, double, int32, int64}"}, {});
+  InstantiationResult result;
+  HasError(
+      InstantiateFunction(fdef, Attrs({{"U", DT_FLOAT}}), GetOpSig, &result),
+      "Attr T is not found from ");
+}
+
+#if 0  // TODO(josh11b): Enable this test once having an extra attr is an error.
+TEST(InstantiateErrors, Too_Many_Attrs) {
+  auto fdef =
+      FDH::Define("nop", {}, {}, {"T:{float, double, int32, int64}"}, {});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, Attrs({{"T", DT_INT32}, {"U", DT_FLOAT}}),
+                               GetOpSig, &result),
+           "Attr U is not found in ");
+}
+#endif
+
+TEST(InstantiateErrors, AttrValue_Value_Placeholder) {
+  auto fdef =
+      FDH::Define("nop", {}, {}, {"T:{float, double, int32, int64}"}, {});
+  InstantiationResult result;
+  HasError(
+      InstantiateFunction(fdef, Attrs({{"T", "$bad"}}), GetOpSig, &result),
+      "AttrValue had value with unexpected type 'placeholder'\n\tfor attr 'T'");
+}
+
+TEST(InstantiateErrors, Unbounded_Attr) {
+  auto fdef = FDH::Define("test", {}, {}, {"T:{float, double, int32, int64}"},
+                          {
+                              {{"a"}, "One", {}, {{"T", "$unknown"}}, {"x"}},
+                          });
+  InstantiationResult result;
+  HasError(
+      InstantiateFunction(fdef, Attrs({{"T", DT_FLOAT}}), GetOpSig, &result),
+      "Failed to bind all placeholders");
+}
+
+TEST(InstantiateErrors, DupArgs) {
+  auto fdef = FDH::Define("test", {"x:float", "x:float"}, {}, {}, {});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Duplicated arg name");
+}
+
+TEST(InstantiateErrors, Dup_Node_Names) {
+  auto fdef = FDH::Define("test", {"x:float"}, {}, {},
+                          {
+                              {{"y"}, "One", {}, {{"T", DT_FLOAT}}},
+                              {{"y"}, "One", {}, {{"T", DT_FLOAT}}},
+                          });
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Duplicated ret name");
+}
+
+TEST(InstantiateErrors, Node_Arg_Notfound) {
+  auto fdef = FDH::Create("test", {"x:float"}, {}, {},
+                          {
+                              {{"y"}, "Add", {"x", "z"}, {{"T", DT_FLOAT}}},
+                          },
+                          {});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "input z is not found");
+}
+
+TEST(InstantiateErrors, Node_Arg_TypeMismatch) {
+  auto fdef = FDH::Define("test", {"x:float"}, {}, {},
+                          {
+                              {{"y"}, "Add", {"x", "x"}, {{"T", DT_INT32}}},
+                          });
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "input x[0] expected type int32 != float, the type of x[0]");
+}
+
+TEST(InstantiateErrors, Node_Arg_ControlMissing) {
+  auto fdef =
+      FDH::Define("test", {"x:float"}, {}, {},
+                  {
+                      {{"y"}, "Add", {"x", "x"}, {{"T", DT_FLOAT}}, {"z"}},
+                  });
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "input[2] == '^z', is not found.");
+}
+
+TEST(InstantiateErrors, FuncRet_Missing) {
+  auto fdef = FDH::Create("test", {}, {"y: float"}, {},
+                          {
+                              {{"x"}, "One", {}, {{"T", DT_FLOAT}}},
+                          },
+                          {});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Return y missing");
+}
+
+TEST(InstantiateErrors, FuncRet_NotFound) {
+  auto fdef = FDH::Create("test", {}, {"y: float"}, {},
+                          {
+                              {{"x"}, "One", {}, {{"T", DT_FLOAT}}},
+                          },
+                          {{"y", "z"}});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Return y -> z is not found");
+}
+
+TEST(InstantiateErrors, FuncRet_NameMismatch) {
+  auto fdef = FDH::Create("test", {}, {"y: float"}, {},
+                          {
+                              {{"x"}, "One", {}, {{"T", DT_FLOAT}}},
+                          },
+                          {{"z", "x:y:0"}});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Return y missing");
+}
+
+// TODO(josh11b): Make this an error.
+// TEST(InstantiateErrors, FuncRet_Extra) {
+//   auto fdef = FDH::Create("test", {}, {"y: float"}, {},
+//                           {
+//                               {{"x"}, "One", {}, {{"T", DT_FLOAT}}},
+//                           },
+//                           {{"y", "x:y:0"}, {"z", "x:y:0"}});
+//   InstantiationResult result;
+//   HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+//            "ret is not found");
+// }
+
+TEST(InstantiateErrors, FuncRet_TypeMismatch) {
+  auto fdef = FDH::Define("test", {}, {"y: float"}, {},
+                          {
+                              {{"y"}, "One", {}, {{"T", DT_DOUBLE}}},
+                          });
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Invalid ret types y : float vs. double\n\tIn function output y");
+}
+
+TEST(InstantiateErrors, TypeList_Missing_Retval_Attr) {
+  auto fdef = FDH::Create(
+      // Name
+      "MySelect",
+      // Args
+      {"x: float"},
+      // Return values
+      {"y: float"},
+      // Attrs
+      {},
+      // Nodes
+      {
+          {{"y"},
+           "Cond",
+           {"x", "x"},
+           {{"tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+            {"cond", FDH::FunctionRef("MyCond2")},
+            {"then_branch", FDH::FunctionRef("MyThen2")},
+            {"else_branch", FDH::FunctionRef("MyElse2")}}},
+      },
+      {{"y", "y:output"}});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "type attr not found: out_types");
+}
+
+TEST(InstantiateErrors, TypeList_Num_Retval_Mismatch) {
+  auto fdef = FDH::Create(
+      // Name
+      "MySelect",
+      // Args
+      {"x: float"},
+      // Return values
+      {"y: float"},
+      // Attrs
+      {},
+      // Nodes
+      {
+          {{"y"},
+           "Cond",
+           {"x", "x"},
+           {{"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+            {"out_types", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+            {"cond", FDH::FunctionRef("MyCond2")},
+            {"then_branch", FDH::FunctionRef("MyThen2")},
+            {"else_branch", FDH::FunctionRef("MyElse2")}}},
+      },
+      {{"y", "y:output"}});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Invalid ret types");
+}
+
+TEST(InstantiateErrors, TypeList_Missing_Arg) {
+  auto fdef = FDH::Create(
+      // Name
+      "MySelect",
+      // Args
+      {"x: float"},
+      // Return values
+      {"y: float"},
+      // Attrs
+      {},
+      // Nodes
+      {
+          {{"y"},
+           "Cond",
+           {"x", "unknown"},
+           {{"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+            {"out_types", DataTypeSlice{DT_FLOAT}},
+            {"cond", FDH::FunctionRef("MyCond2")},
+            {"then_branch", FDH::FunctionRef("MyThen2")},
+            {"else_branch", FDH::FunctionRef("MyElse2")}}},
+      },
+      {{"y", "y:output"}});
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "input unknown is not found");
+}
+
+TEST(InstantiateErrors, TooManyInputs) {
+  auto fdef = FDH::Create(
+      // Name
+      "TooManyInputs",
+      // Inputs
+      {"x: float", "y: float"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = AddN<N=2>(x, y, x)
+       {{"a"}, "AddN", {"x", "y", "x"}, {{"T", DT_FLOAT}, {"N", 2}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Expected input[2] == 'x' to be a control input.");
+}
+
+TEST(InstantiateErrors, TooFewInputs) {
+  auto fdef = FDH::Create(
+      // Name
+      "TooFewInputs",
+      // Inputs
+      {"x: float", "y: float"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = AddN<N=3>(x, y)
+       {{"a"}, "AddN", {"x", "y"}, {{"T", DT_FLOAT}, {"N", 3}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Attempt to access beyond input size: 2 >= 2");
+}
+
+TEST(InstantiateErrors, TooManyInputsFromArray1) {
+  auto fdef = FDH::Create(
+      // Name
+      "TooManyInputsFromArray",
+      // Inputs
+      {"x: float", "y: float"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = _ListToArray(x,y)
+       {{"a"},
+        "_ListToArray",
+        {"x", "y"},
+        {{"N", 2},
+         {"T", DT_FLOAT},
+         {"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}}}},
+       // b = AddN<N=2>(a, y)
+       {{"b"}, "AddN", {"a:output", "y"}, {{"T", DT_FLOAT}, {"N", 2}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Expected input[1] == 'y' to be a control input.");
+}
+
+TEST(InstantiateErrors, TooManyInputsFromArray2) {
+  auto fdef = FDH::Create(
+      // Name
+      "TooManyInputsFromArray",
+      // Inputs
+      {"x: float", "y: float"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = _ListToArray(x,y)
+       {{"a"},
+        "_ListToArray",
+        {"x", "y"},
+        {{"N", 2},
+         {"T", DT_FLOAT},
+         {"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}}}},
+       // b = AddN<N=2>(x, a)
+       {{"b"}, "AddN", {"x", "a:output"}, {{"T", DT_FLOAT}, {"N", 2}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "Input a:output too long for inputs");
+}
+
+TEST(InstantiateErrors, TypeMismatch) {
+  auto fdef = FDH::Create(
+      // Name
+      "TypeMismatch",
+      // Inputs
+      {"x: float", "y: int32"},
+      // Outputs
+      {"z: float"},
+      // Attrs
+      {},
+      // Nodes
+      {// a = AddN<N=2>(x, y)
+       {{"a"}, "AddN", {"x", "y"}, {{"T", DT_FLOAT}, {"N", 3}}}},
+      // Returns
+      {{"z", "a:sum:0"}});
+
+  InstantiationResult result;
+  HasError(InstantiateFunction(fdef, AttrSlice(), GetOpSig, &result),
+           "input inputs[1] expected type float != int32, the type of y[0]");
+}
+
+TEST(FunctionCallFrame, Void_Void) {
+  FunctionCallFrame frame({}, {});
+  TF_EXPECT_OK(frame.SetArgs({}));
+  auto a = test::AsTensor<float>({100});
+  HasError(frame.SetArgs({a}), "Invalid argument");
+  Tensor v;
+  HasError(frame.GetArg(0, &v), "Invalid argument");
+  HasError(frame.SetRetval(0, v), "Invalid argument");
+  std::vector<Tensor> rets;
+  TF_EXPECT_OK(frame.GetRetvals(&rets));
+  EXPECT_EQ(rets.size(), 0);
+}
+
+TEST(FunctionCallFrame, Float_Float_Float) {
+  FunctionCallFrame frame({DT_FLOAT, DT_FLOAT}, {DT_FLOAT});
+  HasError(frame.SetArgs({}), "Invalid argument: Expects 2 arguments");
+  auto a = test::AsTensor<float>({100});
+  auto b = test::AsTensor<float>({200});
+  auto c = test::AsTensor<int64>({300});
+  HasError(frame.SetArgs({a, c}),
+           "Invalid argument: Expects arg[1] to be float");
+  TF_EXPECT_OK(frame.SetArgs({a, b}));
+
+  Tensor v;
+  HasError(frame.GetArg(-1, &v), "Invalid argument");
+  HasError(frame.GetArg(2, &v), "Invalid argument");
+  TF_EXPECT_OK(frame.GetArg(0, &v));
+  test::ExpectTensorEqual<float>(a, v);
+  TF_EXPECT_OK(frame.GetArg(1, &v));
+  test::ExpectTensorEqual<float>(b, v);
+
+  v = test::AsTensor<float>({-100});
+  HasError(frame.SetRetval(-1, v), "Invalid argument");
+  HasError(frame.SetRetval(1, v), "Invalid argument");
+  HasError(frame.SetRetval(0, test::AsTensor<int64>({-100})),
+           "Invalid argument: Expects ret[0] to be float");
+
+  std::vector<Tensor> rets;
+  HasError(frame.GetRetvals(&rets), "does not have value");
+  TF_EXPECT_OK(frame.SetRetval(0, v));
+  HasError(frame.SetRetval(0, v), "has already been set");
+
+  TF_EXPECT_OK(frame.GetRetvals(&rets));
+  EXPECT_EQ(rets.size(), 1);
+  test::ExpectTensorEqual<float>(rets[0], v);
+}
+
+TEST(Canonicalize, Basic) {
+  EXPECT_EQ(Canonicalize("MatMul", Attrs({{"T", DT_FLOAT},
+                                          {"transpose_a", false},
+                                          {"transpose_b", false}})),
+            "MatMul[T=float,transpose_a=false,transpose_b=false]");
+  EXPECT_EQ(Canonicalize("MatMul", Attrs({{"T", DT_FLOAT},
+                                          {"transpose_b", false},
+                                          {"transpose_a", false}})),
+            "MatMul[T=float,transpose_a=false,transpose_b=false]");
+  EXPECT_EQ(Canonicalize("MatMul", Attrs({{"T", DT_DOUBLE},
+                                          {"transpose_b", true},
+                                          {"transpose_a", false}})),
+            "MatMul[T=double,transpose_a=false,transpose_b=true]");
+}
+
+TEST(FunctionLibraryDefinitionTest, Find) {
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+
+  EXPECT_EQ(lib_def.Find("XTimes16"), nullptr);
+
+  auto expect = R"P(
+XTimesTwo[T:{float, double, int32, int64}](x:T) -> (y:T) {
+  two = Const[dtype=int64, value=Tensor<type: int64 shape: [] values: 2>]()
+  scale = Cast[DstT=$T, SrcT=int64](two:output:0)
+  y = Mul[T=$T](x, scale:y:0)
+  return y = y:z:0
+}
+)P";
+  auto found = lib_def.Find("XTimesTwo");
+  ASSERT_NE(found, nullptr);
+  EXPECT_EQ(expect, DebugString(*found));
+}
+
+TEST(FunctionLibraryDefinitionTest, LookUp) {
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+
+  const OpDef* op_def;
+  EXPECT_TRUE(!lib_def.LookUpOpDef("XTimes16", &op_def).ok());
+
+  TF_EXPECT_OK(lib_def.LookUpOpDef("XTimesTwo", &op_def));
+  ASSERT_NE(op_def, nullptr);
+  EXPECT_EQ(op_def->DebugString(),
+            test::function::XTimesTwo().signature().DebugString());
+
+  const OpRegistrationData* op_reg_data;
+  TF_EXPECT_OK(lib_def.LookUp("XTimesTwo", &op_reg_data));
+  ASSERT_NE(op_reg_data, nullptr);
+  // Shape inference function is initialized to UnknownShape.
+  ASSERT_NE(op_reg_data->shape_inference_fn, nullptr);
+}
+
+TEST(FunctionLibraryDefinitionTest, AddFunctionDef) {
+  // Add one function to the proto lib before constructing 'lib_def'.
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+
+  // Add a new function def to the library.
+  TF_EXPECT_OK(lib_def.AddFunctionDef(test::function::WXPlusB()));
+
+  // Test lookup of first function.
+  const OpDef* first;
+  TF_EXPECT_OK(lib_def.LookUpOpDef("XTimesTwo", &first));
+  ASSERT_NE(first, nullptr);
+  EXPECT_EQ(first->DebugString(),
+            test::function::XTimesTwo().signature().DebugString());
+
+  // Test lookup of second function.
+  const OpDef* second;
+  TF_EXPECT_OK(lib_def.LookUpOpDef("WXPlusB", &second));
+  ASSERT_NE(second, nullptr);
+  EXPECT_EQ(second->DebugString(),
+            test::function::WXPlusB().signature().DebugString());
+
+  // Can't add function with same name as existing op
+  FunctionDef fdef = test::function::XTimesTwo();
+  fdef.mutable_signature()->set_name("Add");
+  Status s = lib_def.AddFunctionDef(fdef);
+  EXPECT_FALSE(s.ok());
+  EXPECT_EQ(s.error_message(),
+            "Cannot add function 'Add' because an op with the same name "
+            "already exists.");
+
+  // Already-added functions don't produce error
+  TF_EXPECT_OK(lib_def.AddFunctionDef(test::function::XTimesTwo()));
+  TF_EXPECT_OK(lib_def.AddFunctionDef(test::function::WXPlusB()));
+}
+
+TEST(FunctionLibraryDefinitionTest, AddGradientDef) {
+  // AddGradientDef() doesn't check that functions referenced exist (yet?)
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), FunctionDefLibrary());
+
+  // Test adding a gradient (XTimesFour isn't a valid grad function for
+  // XTimesTwo but that's ok for now)
+  GradientDef grad;
+  grad.set_function_name(test::function::XTimesTwo().signature().name());
+  grad.set_gradient_func(test::function::XTimesFour().signature().name());
+  TF_EXPECT_OK(lib_def.AddGradientDef(grad));
+
+  // Already-added gradients don't produce error
+  TF_EXPECT_OK(lib_def.AddGradientDef(grad));
+
+  // Test that adding a duplicate gradient fails
+  grad.set_gradient_func(test::function::XTimes16().signature().name());
+  Status s = lib_def.AddGradientDef(grad);
+  EXPECT_EQ(s.code(), error::Code::INVALID_ARGUMENT);
+  EXPECT_EQ(s.error_message(),
+            "Cannot assign gradient function 'XTimes16' to 'XTimesTwo' because "
+            "it already has gradient function 'XTimesFour'");
+}
+
+TEST(FunctionLibraryDefinitionTest, AddLibrary) {
+  // Create lib def with single function
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+
+  // Add gradient
+  GradientDef grad;
+  grad.set_function_name(test::function::XTimesTwo().signature().name());
+  grad.set_gradient_func(test::function::XTimesFour().signature().name());
+  TF_EXPECT_OK(lib_def.AddGradientDef(grad));
+
+  // Error if you try to add conflicting function
+  proto.Clear();
+  FunctionDef fdef = test::function::XTimesFour();
+  fdef.mutable_signature()->set_name(
+      test::function::XTimesTwo().signature().name());
+  *proto.add_function() = fdef;
+  FunctionLibraryDefinition lib_def2(OpRegistry::Global(), proto);
+  Status s = lib_def.AddLibrary(lib_def2);
+  EXPECT_EQ(s.code(), error::Code::INVALID_ARGUMENT);
+  EXPECT_EQ(s.error_message(),
+            "Cannot add function 'XTimesTwo' because a different function with "
+            "the same name already exists.");
+
+  // Error if you try to add conflicting gradient
+  proto.Clear();
+  grad.set_gradient_func(test::function::XTimes16().signature().name());
+  *proto.add_gradient() = grad;
+  FunctionLibraryDefinition lib_def3(OpRegistry::Global(), proto);
+  s = lib_def.AddLibrary(lib_def3);
+  EXPECT_EQ(s.code(), error::Code::INVALID_ARGUMENT);
+  EXPECT_EQ(s.error_message(),
+            "Cannot assign gradient function 'XTimes16' to 'XTimesTwo' because "
+            "it already has gradient function 'XTimesFour'");
+
+  // No conflicting functions or gradients OK
+  proto.Clear();
+  *proto.add_function() = test::function::XTimesFour();
+  grad.set_function_name(test::function::XTimes16().signature().name());
+  *proto.add_gradient() = grad;
+  FunctionLibraryDefinition lib_def4(OpRegistry::Global(), proto);
+  TF_EXPECT_OK(lib_def.AddLibrary(lib_def4));
+
+  // OK to add the same functions and gradients twice
+  TF_EXPECT_OK(lib_def.AddLibrary(lib_def));
+}
+
+GradientDef MakeGradDef(const string& f, const string& g) {
+  GradientDef grad;
+  grad.set_function_name(f);
+  grad.set_gradient_func(g);
+  return grad;
+}
+
+TEST(FunctionLibraryDefinitionTest, AddLibrary_Atomic) {
+  // Create lib def containing two functions with equal names
+  FunctionDefLibrary proto;
+  const string x2_name = test::function::XTimesTwo().signature().name();
+  const string x4_name = test::function::XTimesFour().signature().name();
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionDef fdef = test::function::XTimesFour();
+  fdef.mutable_signature()->set_name(x2_name);
+  *proto.add_function() = fdef;
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), FunctionDefLibrary());
+
+  // Try adding the two functions to lib_def
+  Status s = lib_def.AddLibrary(proto);
+  EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
+  EXPECT_EQ(
+      "Cannot add function 'XTimesTwo' because a different function with "
+      "the same name already exists.",
+      s.error_message());
+
+  // Verify that none of the functions are added
+  EXPECT_TRUE(lib_def.Find(x2_name) == nullptr);
+
+  // Fix the name in proto but add two gradient names for it
+  proto.mutable_function(1)->mutable_signature()->set_name(x4_name);
+  *proto.add_gradient() = MakeGradDef(x2_name, x4_name);
+  *proto.add_gradient() = MakeGradDef(x2_name, "SecondGradName");
+
+  // Try adding the library and check that nothing was added
+  s = lib_def.AddLibrary(proto);
+  EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
+  EXPECT_EQ(s.error_message(),
+            "Cannot assign gradient function 'SecondGradName' to 'XTimesTwo' "
+            "because it already has gradient function 'XTimesFour'");
+  EXPECT_TRUE(lib_def.Find(x2_name) == nullptr);
+  EXPECT_EQ(0, lib_def.ToProto().function_size());
+  EXPECT_EQ(0, lib_def.ToProto().gradient_size());
+}
+
+TEST(FunctionLibraryDefinitionTest, AddLibraryDefinition_Atomic_FuncConflict) {
+  const string x2_name = test::function::XTimesTwo().signature().name();
+  const string x4_name = test::function::XTimesFour().signature().name();
+  const string wx_name = test::function::WXPlusB().signature().name();
+
+  // Create FunctionLibraryDefinition with
+  // (func = XTimesTwo, grad = XTimesFour)
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  *proto.add_gradient() = MakeGradDef(x2_name, x4_name);
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+  EXPECT_EQ(1, lib_def.ToProto().function_size());
+  EXPECT_EQ(1, lib_def.ToProto().gradient_size());
+
+  // Create FunctionLibraryDefinition with (func = WXPlusB, grad = XTimesTwo)
+  // and function (name = XTimesTwo, body = XTimeFour)
+  FunctionDefLibrary proto2;
+  *proto2.add_function() = test::function::WXPlusB();
+  *proto2.add_gradient() = MakeGradDef(wx_name, x2_name);
+  *proto2.add_function() = test::function::XTimesFour();
+  proto2.mutable_function(1)->mutable_signature()->set_name(x2_name);
+  FunctionLibraryDefinition lib_def2(OpRegistry::Global(), proto2);
+
+  // Verify that adding lib_def2 will fail because of function conflict
+  // and WXPlusB is not added.
+  Status s = lib_def.AddLibrary(lib_def2);
+  EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
+  EXPECT_EQ(
+      "Cannot add function 'XTimesTwo' because a different function "
+      "with the same name already exists.",
+      s.error_message());
+  EXPECT_TRUE(lib_def.Find(wx_name) == nullptr);
+  EXPECT_EQ(1, lib_def.ToProto().function_size());
+  EXPECT_EQ(1, lib_def.ToProto().gradient_size());
+}
+
+TEST(FunctionLibraryDefinitionTest, AddLibraryDefinition_Atomic_GradConflict) {
+  const string x2_name = test::function::XTimesTwo().signature().name();
+  const string x4_name = test::function::XTimesFour().signature().name();
+  const string wx_name = test::function::WXPlusB().signature().name();
+
+  // Create FunctionLibraryDefinition with
+  // (func = XTimesTwo, grad = XTimesFour)
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  *proto.add_gradient() = MakeGradDef(x2_name, x4_name);
+  FunctionLibraryDefinition lib_def(OpRegistry::Global(), proto);
+  EXPECT_EQ(1, lib_def.ToProto().function_size());
+  EXPECT_EQ(1, lib_def.ToProto().gradient_size());
+
+  // Create FunctionLibraryDefinition with (func = WXPlusB, grad = XTimesTwo)
+  // and (func = XTimesTwo, grad = WXPlusB)
+  FunctionDefLibrary proto2;
+  *proto2.add_function() = test::function::WXPlusB();
+  *proto2.add_gradient() = MakeGradDef(wx_name, x2_name);
+  *proto2.add_function() = test::function::XTimesTwo();
+  *proto2.add_gradient() = MakeGradDef(x2_name, wx_name);
+  FunctionLibraryDefinition lib_def2(OpRegistry::Global(), proto2);
+
+  // Verify that adding lib_def2 will fail because of gradient conflict
+  // and WXPlusB is not added.
+  Status s = lib_def.AddLibrary(lib_def2);
+  EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
+  EXPECT_EQ(
+      "Cannot assign gradient function 'WXPlusB' to 'XTimesTwo'"
+      " because it already has gradient function 'XTimesFour'",
+      s.error_message());
+  EXPECT_TRUE(lib_def.Find(wx_name) == nullptr);
+  EXPECT_EQ(1, lib_def.ToProto().function_size());
+  EXPECT_EQ(1, lib_def.ToProto().gradient_size());
+}
+
+TEST(FunctionLibraryDefinitionTest, ToProto) {
+  FunctionDefLibrary proto1;
+  *proto1.add_function() = test::function::XTimesTwo();
+  *proto1.add_function() = test::function::WXPlusB();
+  FunctionLibraryDefinition lib_def1(OpRegistry::Global(), proto1);
+
+  // Call 'ToProto' and make sure both protos have the same function lib size.
+  FunctionDefLibrary proto2 = lib_def1.ToProto();
+  EXPECT_EQ(proto1.function_size(), proto2.function_size());
+
+  // Initialize 'lib_def2' with proto returned by 'ToProto' call.
+  FunctionLibraryDefinition lib_def2(OpRegistry::Global(), proto2);
+
+  // Test that the first function exists in both libraries.
+  const OpDef *f1, *f2, *f3, *f4;
+  TF_EXPECT_OK(lib_def1.LookUpOpDef("XTimesTwo", &f1));
+  TF_EXPECT_OK(lib_def2.LookUpOpDef("XTimesTwo", &f2));
+  EXPECT_EQ(f1->DebugString(), f2->DebugString());
+
+  // Test that the second function exists in both libraries.
+  TF_EXPECT_OK(lib_def1.LookUpOpDef("WXPlusB", &f3));
+  TF_EXPECT_OK(lib_def2.LookUpOpDef("WXPlusB", &f4));
+  EXPECT_EQ(f3->DebugString(), f4->DebugString());
+}
+
+TEST(FunctionLibraryDefinitionTest, GetAttr_FuncNoAttr) {
+  FunctionDefLibrary proto;
+  *proto.add_function() = test::function::XTimesTwo();
+  FunctionLibraryDefinition lib(OpRegistry::Global(), proto);
+
+  NodeDef ndef;
+  bool annotation;
+
+  // Not a function.
+  ndef.set_op("Matmul");
+  EXPECT_FALSE(lib.GetAttr(ndef, "annotation", &annotation).ok());
+
+  // A function. No attr defined.
+  ndef.set_op("XTimesTwo");
+  EXPECT_FALSE(lib.GetAttr(ndef, "annotation", &annotation).ok());
+
+  // ndef defines the attr. But we don't care.
+  AddNodeAttr("annotation", true, &ndef);
+  EXPECT_FALSE(lib.GetAttr(ndef, "annotation", &annotation).ok());
+}
+
+template <typename T>
+void SetAttrValue(FunctionDef* fdef, const string& attr, const T& value) {
+  AttrValue attr_value;
+  SetAttrValue(value, &attr_value);
+  fdef->mutable_attr()->insert({attr, attr_value});
+}
+
+TEST(FunctionLibraryDefinitionTest, GetAttr_FuncWithAttr) {
+  FunctionDefLibrary proto;
+  auto fdef = proto.add_function();
+  *fdef = test::function::XTimesTwo();
+  SetAttrValue(fdef, "annotation", true);
+  SetAttrValue(fdef, "options", "some string data");
+  FunctionLibraryDefinition lib(OpRegistry::Global(), proto);
+
+  NodeDef ndef;
+  bool annotation;
+
+  // A function. No attr defined in ndef.
+  ndef.set_op("XTimesTwo");
+  TF_EXPECT_OK(lib.GetAttr(ndef, "annotation", &annotation));
+  EXPECT_EQ(annotation, true);
+
+  string str;
+  TF_EXPECT_OK(lib.GetAttr(ndef, "options", &str));
+  EXPECT_EQ(str, "some string data");
+}
+
+TEST(FunctionLibraryDefinitionTest, GetAttr_Gradient) {
+  FunctionDefLibrary proto;
+  auto fdef = proto.add_function();
+  *fdef = test::function::XTimesTwo();
+  SetAttrValue(fdef, "annotation", true);
+  *fdef = test::function::WXPlusB();
+  SetAttrValue(fdef, "annotation", false);
+  auto func_grad = proto.add_gradient();
+  func_grad->set_function_name("XTimesTwo");
+  func_grad->set_gradient_func("WXPlusB");
+  FunctionLibraryDefinition lib(OpRegistry::Global(), proto);
+
+  NodeDef ndef;
+  ndef.set_op(FunctionLibraryDefinition::kGradientOp);
+
+  bool annotation;
+  EXPECT_FALSE(lib.GetAttr(ndef, "annotation", &annotation).ok());
+
+  NameAttrList nal;
+  nal.set_name("XTimesTwo");
+  AddNodeAttr(FunctionLibraryDefinition::kFuncAttr, nal, &ndef);
+  TF_EXPECT_OK(lib.GetAttr(ndef, "annotation", &annotation));
+  EXPECT_EQ(annotation, false);  // XTimesTwo's gradient is WXPlusB.
+
+  nal.set_name("WXPlusB");
+  ndef.clear_attr();
+  AddNodeAttr(FunctionLibraryDefinition::kFuncAttr, nal, &ndef);
+  TF_EXPECT_OK(lib.GetAttr(ndef, "annotation", &annotation));
+  EXPECT_EQ(annotation, false);  // WXPlusB has no custom gradient.
+}
+
+// TODO(skyewm): this could be more thorough
+TEST(FunctionDefsEqualTest, TestFunctionDefsEqual) {
+  // Equal functions
+  const FunctionDef fdef1 = test::function::XTimesTwo();
+  FunctionDef fdef2 = test::function::XTimesTwo();
+  uint64 hash1 = FunctionDefHash(fdef1);
+  EXPECT_TRUE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_EQ(hash1, FunctionDefHash(fdef2));
+
+  // Different functions
+  fdef2 = test::function::XTimesFour();
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Different signatures
+  fdef2 = test::function::XTimesTwo();
+  fdef2.mutable_signature()->mutable_input_arg(0)->set_name("foo");
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Descriptions must be equal
+  fdef2 = test::function::XTimesTwo();
+  fdef2.mutable_signature()->mutable_input_arg(0)->set_description("foo");
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Different NodeDefs
+  fdef2 = test::function::XTimesTwo();
+  NodeDef* ndef = fdef2.add_node_def();
+  *ndef = fdef2.node_def(0);
+  ndef->set_name("new_name");
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Different return values
+  fdef2 = test::function::XTimesTwo();
+  (*fdef2.mutable_ret())["y"] = "y:z:1";  // originally is "y:z:0"
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Different attributes
+  fdef2 = test::function::XTimesTwo();
+  SetAttrValue(&fdef2, "ExtraAttr", true);
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+  EXPECT_NE(hash1, FunctionDefHash(fdef2));
+
+  // Multiple equivalent attributes; the two functions should be equal.
+  fdef2 = test::function::XTimesTwo();
+  FunctionDef fdef3 = test::function::XTimesTwo();
+  SetAttrValue(&fdef2, "Foo", true);
+  SetAttrValue(&fdef3, "Foo", true);
+  SetAttrValue(&fdef2, "Bar", 123);
+  SetAttrValue(&fdef3, "Bar", 123);
+  SetAttrValue(&fdef2, "Baz", "abc");
+  SetAttrValue(&fdef3, "Baz", "abc");
+  EXPECT_TRUE(FunctionDefsEqual(fdef2, fdef3));
+  EXPECT_EQ(FunctionDefHash(fdef2), FunctionDefHash(fdef3));
+}
+
+}  // end namespace
+}  // end namespace tensorflow

function_testlib.cc

@@ -0,0 +1,204 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/function_testlib.h"
+
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace test {
+namespace function {
+
+typedef FunctionDefHelper FDH;
+
+GraphDef GDef(gtl::ArraySlice<NodeDef> nodes,
+              gtl::ArraySlice<FunctionDef> funcs) {
+  GraphDef g;
+  VersionDef* versions = g.mutable_versions();
+  versions->set_producer(TF_GRAPH_DEF_VERSION);
+  versions->set_min_consumer(TF_GRAPH_DEF_VERSION_MIN_CONSUMER);
+  for (const auto& n : nodes) {
+    *(g.add_node()) = n;
+  }
+  auto lib = g.mutable_library();
+  for (const auto& f : funcs) {
+    *(lib->add_function()) = f;
+  }
+  return g;
+}
+
+// Helper to construct a NodeDef.
+NodeDef NDef(const string& name, const string& op,
+             gtl::ArraySlice<string> inputs,
+             gtl::ArraySlice<std::pair<string, FDH::AttrValueWrapper>> attrs,
+             const string& device) {
+  NodeDef n;
+  n.set_name(name);
+  n.set_op(op);
+  for (const auto& in : inputs) n.add_input(in);
+  n.set_device(device);
+  for (auto na : attrs) n.mutable_attr()->insert({na.first, na.second.proto});
+  return n;
+}
+
+FunctionDef NonZero() {
+  return FDH::Define(
+      // Name
+      "NonZero",
+      // Args
+      {"x:T"},
+      // Return values
+      {"y:T"},
+      // Attr def
+      {"T:{float, double, int32, int64, string}"},
+      // Nodes
+      {
+          {{"y"}, "Identity", {"x"}, {{"T", "$T"}}},
+      });
+}
+
+FunctionDef XTimesTwo() {
+  const Tensor kTwo = test::AsScalar<int64>(2);
+  return FDH::Define(
+      // Name
+      "XTimesTwo",
+      // Args
+      {"x: T"},
+      // Return values
+      {"y: T"},
+      // Attr def
+      {"T: {float, double, int32, int64}"},
+      // Nodes
+      {
+          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
+          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
+          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
+      });
+}
+
+FunctionDef XTimesTwoInt32() {
+  const Tensor kTwo = test::AsScalar<int64>(2);
+  return FDH::Define(
+      // Name
+      "XTimesTwoInt32",
+      // Args
+      {"x: int32"},
+      // Return values
+      {"y: int32"}, {},
+      // Nodes
+      {
+          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
+          {{"scale"},
+           "Cast",
+           {"two"},
+           {{"SrcT", DT_INT64}, {"DstT", DT_INT32}}},
+          {{"y"}, "Mul", {"x", "scale"}, {{"T", DT_INT32}}},
+      });
+}
+
+FunctionDef XTimesFour() {
+  return FDH::Create(
+      // Name
+      "XTimesFour",
+      // Args
+      {"x: T"},
+      // Return values
+      {"y: T"},
+      // Attr def
+      {"T: {float, double, int32, int64}"},
+      // Nodes
+      {
+          {{"x2"}, "XTimesTwo", {"x"}, {{"T", "$T"}}},
+          {{"y"}, "XTimesTwo", {"x2:y:0"}, {{"T", "$T"}}},
+      },
+      {{"y", "y:y:0"}});
+}
+
+FunctionDef XTimes16() {
+  return FDH::Create(
+      // Name
+      "XTimes16",
+      // Args
+      {"x: T"},
+      // Return values
+      {"y: T"},
+      // Attr def
+      {"T: {float, double, int32, int64}"},
+      // Nodes
+      {
+          {{"x4"}, "XTimesFour", {"x"}, {{"T", "$T"}}},
+          {{"y"}, "XTimesFour", {"x4:y:0"}, {{"T", "$T"}}},
+      },
+      {{"y", "y:y:0"}});
+}
+
+FunctionDef WXPlusB(){return FDH::Define(
+    // Name
+    "WXPlusB",
+    // Args
+    {"w: T", "x: T", "b: T"},
+    // Return values
+    {"y: T"},
+    // Attr def
+    {"T: {float, double}"},
+    // Nodes
+    {
+      {{"mm"},
+       "MatMul",
+       {"w", "x"},
+       {
+           {"T", "$T"}, {"transpose_a", false}, {"transpose_b", false},
+#ifdef INTEL_MKL
+       }},
+#else
+         {"_kernel", "eigen"}}},
+#endif
+      {
+        {"y"}, "Add", {"mm", "b"}, {
+          { "T", "$T" }
+        }
+      }
+    });
+}
+
+FunctionDef Swap() {
+  return FDH::Define(
+      // Name
+      "Swap",
+      // Args
+      {"i0: T", "i1: T"},
+      // Return values
+      {"o0: T", "o1: T"},
+      // Attr def
+      {"T: {float, double}"},
+      // Nodes
+      {{{"o0"}, "Identity", {"i1"}, {{"T", "$T"}}},
+       {{"o1"}, "Identity", {"i0"}, {{"T", "$T"}}}});
+}
+
+void FunctionTestSchedClosure(std::function<void()> fn) {
+  static thread::ThreadPool* w =
+      new thread::ThreadPool(Env::Default(), "Test", 8);
+  w->Schedule(std::move(fn));
+}
+
+}  // end namespace function
+}  // end namespace test
+}  // end namespace tensorflow

function_testlib.h

@@ -0,0 +1,90 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_
+#define TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_
+
+#include <string>
+
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace test {
+namespace function {
+
+// A helper class to make AttrSlice from initializer lists
+class Attrs {
+ public:
+  Attrs(const std::initializer_list<  // NOLINT(runtime/explicit)
+        std::pair<string, FunctionDefHelper::AttrValueWrapper>>& attrs) {
+    for (const auto& aval : attrs) {
+      map_.insert({aval.first, aval.second.proto});
+    }
+  }
+
+  operator AttrSlice() { return AttrSlice(&map_); }  // NOLINT(runtime/explicit)
+
+ private:
+  AttrValueMap map_;
+};
+
+// Helper to construct a NodeDef.
+NodeDef NDef(
+    const string& name, const string& op, gtl::ArraySlice<string> inputs,
+    gtl::ArraySlice<std::pair<string, FunctionDefHelper::AttrValueWrapper>>
+        attrs = {},
+    const string& device = "");
+
+// Helper to construct a GraphDef proto.
+GraphDef GDef(gtl::ArraySlice<NodeDef> nodes,
+              gtl::ArraySlice<FunctionDef> funcs = {});
+
+// For testing convenience, we provide a few simple functions that can
+// be easily executed and tested.
+
+// x:T -> x * 2.
+FunctionDef XTimesTwo();
+
+// x:T -> x * 2, where x is int32.
+FunctionDef XTimesTwoInt32();
+
+// x:T -> (x * 2) * 2.
+FunctionDef XTimesFour();
+
+// x:T -> ((x * 2) * 2) * 2.
+FunctionDef XTimes16();
+
+// w:T, x:T, b:T -> MatMul(w, x) + b
+FunctionDef WXPlusB();
+
+// x:T -> x:T, T is a type which we automatically converts to a bool.
+FunctionDef NonZero();
+
+// x:T, y:T -> y:T, x:T
+FunctionDef Swap();
+
+void FunctionTestSchedClosure(std::function<void()> fn);
+
+}  // end namespace function
+}  // end namespace test
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_

graph.proto

@@ -0,0 +1,56 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "GraphProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/node_def.proto";
+import "tensorflow/core/framework/function.proto";
+import "tensorflow/core/framework/versions.proto";
+
+// Represents the graph of operations
+message GraphDef {
+  repeated NodeDef node = 1;
+
+  // Compatibility versions of the graph.  See core/public/version.h for version
+  // history.  The GraphDef version is distinct from the TensorFlow version, and
+  // each release of TensorFlow will support a range of GraphDef versions.
+  VersionDef versions = 4;
+
+  // Deprecated single version field; use versions above instead.  Since all
+  // GraphDef changes before "versions" was introduced were forward
+  // compatible, this field is entirely ignored.
+  int32 version = 3 [deprecated = true];
+
+  // EXPERIMENTAL. DO NOT USE OR DEPEND ON THIS YET.
+  //
+  // "library" provides user-defined functions.
+  //
+  // Naming:
+  //   * library.function.name are in a flat namespace.
+  //     NOTE: We may need to change it to be hierarchical to support
+  //     different orgs. E.g.,
+  //     { "/google/nn", { ... }},
+  //     { "/google/vision", { ... }}
+  //     { "/org_foo/module_bar", { ... }}
+  //     map<string, FunctionDefLib> named_lib;
+  //   * If node[i].op is the name of one function in "library",
+  //     node[i] is deemed as a function call. Otherwise, node[i].op
+  //     must be a primitive operation supported by the runtime.
+  //
+  //
+  // Function call semantics:
+  //
+  //   * The callee may start execution as soon as some of its inputs
+  //     are ready. The caller may want to use Tuple() mechanism to
+  //     ensure all inputs are ready in the same time.
+  //
+  //   * The consumer of return values may start executing as soon as
+  //     the return values the consumer depends on are ready.  The
+  //     consumer may want to use Tuple() mechanism to ensure the
+  //     consumer does not start until all return values of the callee
+  //     function are ready.
+  FunctionDefLibrary library = 2;
+};

graph_def_util.cc

@@ -0,0 +1,218 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/graph_def_util.h"
+
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op_def_util.h"
+#include "tensorflow/core/framework/versions.pb_text.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+
+namespace tensorflow {
+
+string SummarizeGraphDef(const GraphDef& graph_def) {
+  string ret;
+  strings::StrAppend(&ret, "versions = ",
+                     ProtoShortDebugString(graph_def.versions()), ";\n");
+  for (const NodeDef& node : graph_def.node()) {
+    strings::StrAppend(&ret, SummarizeNodeDef(node), ";\n");
+  }
+  return ret;
+}
+
+Status ValidateExternalGraphDefSyntax(const GraphDef& graph_def) {
+  for (const NodeDef& node : graph_def.node()) {
+    TF_RETURN_IF_ERROR(ValidateExternalNodeDefSyntax(node));
+  }
+  return Status::OK();
+}
+
+Status AddDefaultAttrsToGraphDef(GraphDef* graph_def,
+                                 const OpRegistryInterface& op_registry,
+                                 int node_offset) {
+  if (node_offset > graph_def->node_size()) {
+    return errors::InvalidArgument(
+        "Tried to add default attrs to GraphDef "
+        "starting at offset ",
+        node_offset, " with total nodes in graph: ", graph_def->node_size());
+  }
+
+  for (int i = node_offset; i < graph_def->node_size(); ++i) {
+    NodeDef* node_def = graph_def->mutable_node(i);
+    const OpDef* op_def;
+    TF_RETURN_IF_ERROR(op_registry.LookUpOpDef(node_def->op(), &op_def));
+    AddDefaultsToNodeDef(*op_def, node_def);
+  }
+
+  return Status::OK();
+}
+
+static Status RemoveNewDefaultAttrsFromNodeDef(
+    NodeDef* node_def, const OpRegistryInterface& consumer_op_registry,
+    const OpRegistryInterface& producer_op_registry,
+    std::set<std::pair<string, string>>* op_attr_removed) {
+  const OpDef* producer_op_def;
+  const OpDef* consumer_op_def;
+  TF_RETURN_IF_ERROR(
+      producer_op_registry.LookUpOpDef(node_def->op(), &producer_op_def));
+  TF_RETURN_IF_ERROR(
+      consumer_op_registry.LookUpOpDef(node_def->op(), &consumer_op_def));
+
+  std::vector<string> to_remove;
+  for (const auto& attr : node_def->attr()) {
+    // If the attr is not in consumer_op_def and doesn't start with '_'...
+    if (!StringPiece(attr.first).starts_with("_") &&
+        FindAttr(attr.first, *consumer_op_def) == nullptr) {
+      const OpDef::AttrDef* producer_attr_def =
+          FindAttr(attr.first, *producer_op_def);
+      if (producer_attr_def == nullptr) {
+        return errors::InvalidArgument(
+            "Attr '", attr.first, "' missing in producer's OpDef: ",
+            SummarizeOpDef(*producer_op_def), " but found in node: ",
+            SummarizeNodeDef(*node_def));
+      }
+      // ...and it has the same value as the default in producer,
+      if (producer_attr_def->has_default_value() &&
+          AreAttrValuesEqual(producer_attr_def->default_value(), attr.second)) {
+        // then we will remove it below.
+        to_remove.emplace_back(attr.first);
+      }
+    }
+  }
+  // We separate identifying which attrs should be removed from
+  // actually removing them to avoid invalidating the loop iterators
+  // above.
+  for (const string& attr_name : to_remove) {
+    node_def->mutable_attr()->erase(attr_name);
+    if (op_attr_removed != nullptr) {
+      op_attr_removed->insert(std::make_pair(node_def->op(), attr_name));
+    }
+  }
+
+  return Status::OK();
+}
+
+static bool IsFunction(const GraphDef& graph_def, const string& op_name) {
+  for (const auto& func_def : graph_def.library().function()) {
+    if (op_name == func_def.signature().name()) return true;
+  }
+  return false;
+}
+
+Status RemoveNewDefaultAttrsFromGraphDef(
+    GraphDef* graph_def, const OpRegistryInterface& consumer_op_registry,
+    const OpRegistryInterface& producer_op_registry,
+    std::set<std::pair<string, string>>* op_attr_removed) {
+  // TODO(joshL): Make IsFunction() faster by collecting the names of
+  // all functions as a preprocessing step.
+  for (int n = 0; n < graph_def->node_size(); ++n) {
+    NodeDef* node_def = graph_def->mutable_node(n);
+    if (!IsFunction(*graph_def, node_def->op())) {
+      TF_RETURN_IF_ERROR(RemoveNewDefaultAttrsFromNodeDef(
+          node_def, consumer_op_registry, producer_op_registry,
+          op_attr_removed));
+    }
+  }
+  for (int f = 0; f < graph_def->library().function_size(); ++f) {
+    FunctionDef* func_def = graph_def->mutable_library()->mutable_function(f);
+    for (int n = 0; n < func_def->node_def_size(); ++n) {
+      NodeDef* node_def = func_def->mutable_node_def(n);
+      if (!IsFunction(*graph_def, node_def->op())) {
+        // TODO(josh11b): Better handling of attrs with placeholder values.
+        TF_RETURN_IF_ERROR(RemoveNewDefaultAttrsFromNodeDef(
+            node_def, consumer_op_registry, producer_op_registry,
+            op_attr_removed));
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+void OpsUsedByGraph(const GraphDef& graph_def,
+                    std::set<string>* ops_used_in_graph) {
+  // Map function names to definitions.
+  std::unordered_map<string, const FunctionDef*> name_to_function;
+  for (const auto& function : graph_def.library().function()) {
+    name_to_function.insert(
+        std::make_pair(function.signature().name(), &function));
+  }
+
+  // Collect the sorted list of op names.  Since functions can reference
+  // functions, we need a recursive traversal.
+  std::set<string> used_ops;  // Includes both primitive ops and functions
+  std::vector<const FunctionDef*> functions_to_process;  // A subset of used_ops
+  // Collect the logic to mark an op in a lambda; it'll be used twice below.
+  const auto mark_op_as_used = [&used_ops, &functions_to_process,
+                                &name_to_function](const string& op) {
+    if (used_ops.insert(op).second) {
+      // If it's a function, we'll need to process further
+      const auto it = name_to_function.find(op);
+      if (it != name_to_function.end()) {
+        functions_to_process.push_back(it->second);
+      }
+    }
+  };
+  for (const auto& node : graph_def.node()) {
+    mark_op_as_used(node.op());
+  }
+  while (!functions_to_process.empty()) {
+    const FunctionDef* fun = functions_to_process.back();
+    functions_to_process.pop_back();
+    for (const auto& node : fun->node_def()) {
+      mark_op_as_used(node.op());
+    }
+  }
+
+  // Filter out function names to produce output.
+  // TODO(josh11b): Change the above code to produce this directly.
+  ops_used_in_graph->clear();
+  for (const string& op_name : used_ops) {
+    if (name_to_function.find(op_name) == name_to_function.end()) {
+      ops_used_in_graph->insert(op_name);
+    }
+  }
+}
+
+Status StrippedOpListForGraph(const GraphDef& graph_def,
+                              const OpRegistryInterface& op_registry,
+                              OpList* stripped_op_list) {
+  std::set<string> used_ops;
+  OpsUsedByGraph(graph_def, &used_ops);
+
+  // Build the stripped op list in sorted order, ignoring functions.
+  stripped_op_list->clear_op();
+  for (const string& op_name : used_ops) {
+    const OpDef* op_def;
+    TF_RETURN_IF_ERROR(op_registry.LookUpOpDef(op_name, &op_def));
+    OpDef* stripped_op = stripped_op_list->add_op();
+    stripped_op->CopyFrom(*op_def);
+    RemoveDescriptionsFromOpDef(stripped_op);
+  }
+  return Status::OK();
+}
+
+}  // namespace tensorflow

graph_def_util.h

@@ -0,0 +1,115 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_GRAPH_DEF_UTIL_H_
+#define TENSORFLOW_FRAMEWORK_GRAPH_DEF_UTIL_H_
+
+#include <set>
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+// Forward declare proto so that it's symbols can be removed from .so exports
+class GraphDef;
+
+// Produce a human-readable version of a GraphDef that is more concise
+// than a text-format proto.
+string SummarizeGraphDef(const GraphDef& graph_def);
+
+// Validates the syntax of a GraphDef provided externally.
+//
+// The following is an EBNF-style syntax for GraphDef objects. Note that
+// Node objects are actually specified as tensorflow::NodeDef protocol buffers,
+// which contain many other fields that are not (currently) validated.
+//
+// Graph        = Node *
+// Node         = NodeName, Inputs
+// Inputs       = ( DataInput * ), ( ControlInput * )
+// DataInput    = NodeName, ( ":", [1-9], [0-9] * ) ?
+// ControlInput = "^", NodeName
+// NodeName     = [A-Za-z0-9.], [A-Za-z0-9_./] *
+Status ValidateExternalGraphDefSyntax(const GraphDef& graph_def);
+
+// Adds default attributes to NodeDefs in 'graph_def' starting
+// from the 'node_offset' node in 'graph_def'.
+//
+// Default attributes are defined by 'op_registry'.
+//
+// Returns OK on success, an error if 'graph_def' has a NodeDef
+// that cannot be found in 'op_registry'.
+//
+// REQUIRES: 'graph_def' and 'op_registry' are not nullptr.
+Status AddDefaultAttrsToGraphDef(GraphDef* graph_def,
+                                 const OpRegistryInterface& op_registry,
+                                 int node_offset);
+
+// Remove attrs from 'graph_def' that have the default value according
+// to 'producer_op_registry', but don't exist according to
+// 'consumer_op_registry'. This can allow 'graph_def' to run on the
+// consumer even if consumer was built at an earlier CL (before an
+// attr with a default was added). Note that this will not affect
+// attrs with non-default values, so you must run a
+// ValidateGraphDef...() function to see if the result is in fact
+// compatible. If not nullptr, the op/attr pairs that were removed
+// are added to '*op_attr_removed'.
+//
+// Expected usage, for a producer that wants to prepare a graph for
+// a consumer:
+// // For each consumer, update 'graph_def':
+//   OpListOpRegistry consumer_op_registry(consumer_server_op_list);
+//   std::unordered_set<std::pair<string, string>> op_attr_removed;
+//   TF_RETURN_IF_ERROR(RemoveNewDefaultAttrsFromGraphDef(
+//       &graph_def, consumer_op_registry, *OpRegistry::Global(),
+//       &op_attr_removed));
+// // Validate that each consumer can understand the resulting 'graph_def'
+//   TF_RETURN_IF_ERROR(graph::ValidateGraphDefAgainstOpRegistry(
+//       graph_def, consumer_op_registry));
+// // Consumer can use 'graph_def', and 'op_attr_removed' summarizes
+// // what changes had to be made to 'graph_def' for it to work.
+//
+// Expected usage, for a consumer that has a graph and a
+// (optionally-stripped) op_list from a producer (say from a call to
+// StrippedOpListForGraph(), or in the MetaGraphDef):
+//   OpListOpRegistry producer_op_registry(producer_stripped_op_list);
+//   TF_RETURN_IF_ERROR(RemoveNewDefaultAttrsFromGraphDef(
+//       &graph_def, *OpRegistry::Global(), producer_op_registry, nullptr));
+Status RemoveNewDefaultAttrsFromGraphDef(
+    GraphDef* graph_def, const OpRegistryInterface& consumer_op_registry,
+    const OpRegistryInterface& producer_op_registry,
+    std::set<std::pair<string, string>>* op_attr_removed);
+
+// Two functions that collect the ops used by a graph.
+//
+// This returns the ops used as a set of strings.
+void OpsUsedByGraph(const GraphDef& graph_def,
+                    std::set<string>* ops_used_in_graph);
+
+// This function computes the stripped_op_list field of MetaGraphDef
+// and similar protos.  The op_registry should contain the ops used to
+// produce graph_def.  The resulting stripped_op_list can be
+// communicated from the producer to the consumer, which can use
+// RemoveNewDefaultAttrsFromGraphDef() to improve forwards compatibility
+// (using an OpListOpRegistry as indicated in the example above).
+//
+// Most users will pass *OpRegistry::Global() for op_registry to strip against
+// the list of ops registered in this process.
+Status StrippedOpListForGraph(const GraphDef& graph_def,
+                              const OpRegistryInterface& op_registry,
+                              OpList* stripped_op_list);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_GRAPH_DEF_UTIL_H_

graph_def_util_test.cc

@@ -0,0 +1,321 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/graph_def_util.h"
+
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/equal_graph_def.h"
+
+namespace tensorflow {
+namespace {
+
+Status FinalizeOpDef(const OpDefBuilder& b, OpDef* op_def) {
+  OpRegistrationData op_reg_data;
+  const Status s = b.Finalize(&op_reg_data);
+  *op_def = op_reg_data.op_def;
+  return s;
+}
+
+// Producer and consumer have default for an attr -> graph unchanged.
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, NoChangeWithDefault) {
+  OpList op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("NoChangeWithDefault").Attr("a: int = 12"),
+                    op_list.add_op()));
+  OpListOpRegistry registry(&op_list);
+
+  GraphDef graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("ncwd", "NoChangeWithDefault", &registry)
+                   .Finalize(graph_def.add_node()));
+  GraphDef expected_graph_def = graph_def;
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(RemoveNewDefaultAttrsFromGraphDef(&graph_def, registry, registry,
+                                                 &op_attr_removed));
+
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, graph_def);
+  EXPECT_TRUE(op_attr_removed.empty());
+}
+
+// Producer and consumer both have an attr -> graph unchanged.
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, NoChangeNoDefault) {
+  OpList op_list;
+  TF_ASSERT_OK(FinalizeOpDef(OpDefBuilder("NoChangeNoDefault").Attr("a: int"),
+                             op_list.add_op()));
+  OpListOpRegistry registry(&op_list);
+
+  GraphDef graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("ncnd", "NoChangeNoDefault", &registry)
+                   .Attr("a", 42)
+                   .Finalize(graph_def.add_node()));
+  GraphDef expected_graph_def = graph_def;
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(RemoveNewDefaultAttrsFromGraphDef(&graph_def, registry, registry,
+                                                 &op_attr_removed));
+
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, graph_def);
+  EXPECT_TRUE(op_attr_removed.empty());
+}
+
+// Producer has default for an attr that the consumer does not know
+// about, and the produced graph has the default value for the attr ->
+// attr removed from graph (and so able to be consumed).
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, UsesDefault) {
+  OpList consumer_op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("UsesDefault"), consumer_op_list.add_op()));
+  OpListOpRegistry consumer_registry(&consumer_op_list);
+
+  OpList producer_op_list;
+  TF_ASSERT_OK(FinalizeOpDef(OpDefBuilder("UsesDefault").Attr("a: int = 17"),
+                             producer_op_list.add_op()));
+  OpListOpRegistry producer_registry(&producer_op_list);
+
+  GraphDef produced_graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("uses_default", "UsesDefault", &producer_registry)
+                   .Finalize(produced_graph_def.add_node()));
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(
+      RemoveNewDefaultAttrsFromGraphDef(&produced_graph_def, consumer_registry,
+                                        producer_registry, &op_attr_removed));
+
+  GraphDef expected_graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("uses_default", "UsesDefault", &consumer_registry)
+                   .Finalize(expected_graph_def.add_node()));
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, produced_graph_def);
+
+  std::set<std::pair<string, string>> expected_removed({{"UsesDefault", "a"}});
+  EXPECT_EQ(expected_removed, op_attr_removed);
+}
+
+// Producer has default for an attr that the consumer does not know
+// about, graph sets the attr to a value different from the default ->
+// graph unchanged (but not able to be consumed by consumer).
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, ChangedFromDefault) {
+  OpList consumer_op_list;
+  TF_ASSERT_OK(FinalizeOpDef(OpDefBuilder("ChangedFromDefault"),
+                             consumer_op_list.add_op()));
+  OpListOpRegistry consumer_registry(&consumer_op_list);
+
+  OpList producer_op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("ChangedFromDefault").Attr("a: int = 17"),
+                    producer_op_list.add_op()));
+  OpListOpRegistry producer_registry(&producer_op_list);
+
+  GraphDef produced_graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("changed_from_default", "ChangedFromDefault",
+                              &producer_registry)
+                   .Attr("a", 9)
+                   .Finalize(produced_graph_def.add_node()));
+  GraphDef expected_graph_def = produced_graph_def;
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(
+      RemoveNewDefaultAttrsFromGraphDef(&produced_graph_def, consumer_registry,
+                                        producer_registry, &op_attr_removed));
+
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, produced_graph_def);
+  EXPECT_TRUE(op_attr_removed.empty());
+}
+
+// Attrs starting with underscores should not be removed.
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, UnderscoreAttrs) {
+  OpList consumer_op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("Underscore"), consumer_op_list.add_op()));
+  OpListOpRegistry consumer_registry(&consumer_op_list);
+
+  OpList producer_op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("Underscore"), producer_op_list.add_op()));
+  // Add the _underscore attr manually since OpDefBuilder would complain
+  OpDef::AttrDef* attr = producer_op_list.mutable_op(0)->add_attr();
+  attr->set_name("_underscore");
+  attr->set_type("int");
+  attr->mutable_default_value()->set_i(17);
+  OpListOpRegistry producer_registry(&producer_op_list);
+
+  GraphDef produced_graph_def;
+  TF_ASSERT_OK(NodeDefBuilder("node", "Underscore", &producer_registry)
+                   .Attr("_underscore", 17)
+                   .Finalize(produced_graph_def.add_node()));
+  GraphDef expected_graph_def = produced_graph_def;
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(
+      RemoveNewDefaultAttrsFromGraphDef(&produced_graph_def, consumer_registry,
+                                        producer_registry, &op_attr_removed));
+
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, produced_graph_def);
+  EXPECT_EQ(op_attr_removed.size(), 0);
+}
+
+TEST(RemoveNewDefaultAttrsFromGraphDefTest, HasFunction) {
+  OpList consumer_op_list;
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("UsesDefault"), consumer_op_list.add_op()));
+  TF_ASSERT_OK(FinalizeOpDef(OpDefBuilder("ChangedFromDefault"),
+                             consumer_op_list.add_op()));
+  OpListOpRegistry consumer_registry(&consumer_op_list);
+
+  OpList producer_op_list;
+  TF_ASSERT_OK(FinalizeOpDef(OpDefBuilder("UsesDefault").Attr("a: int = 17"),
+                             producer_op_list.add_op()));
+  TF_ASSERT_OK(
+      FinalizeOpDef(OpDefBuilder("ChangedFromDefault").Attr("a: int = 17"),
+                    producer_op_list.add_op()));
+  OpListOpRegistry producer_registry(&producer_op_list);
+
+  GraphDef produced_graph_def;
+  *produced_graph_def.mutable_library()->add_function() =
+      FunctionDefHelper::Create(
+          "my_func", {}, {}, {},
+          {{{"x"}, "UsesDefault", {}, {{"a", 17}}},
+           {{"y"}, "ChangedFromDefault", {}, {{"a", 99}}}},
+          {});
+  OpList function_op_list;
+  *function_op_list.add_op() =
+      produced_graph_def.library().function(0).signature();
+  OpListOpRegistry function_registry(&function_op_list);
+  TF_ASSERT_OK(NodeDefBuilder("call_func", "my_func", &function_registry)
+                   .Finalize(produced_graph_def.add_node()));
+
+  std::set<std::pair<string, string>> op_attr_removed;
+  TF_ASSERT_OK(
+      RemoveNewDefaultAttrsFromGraphDef(&produced_graph_def, consumer_registry,
+                                        producer_registry, &op_attr_removed));
+
+  GraphDef expected_graph_def;
+  *expected_graph_def.mutable_library()->add_function() =
+      FunctionDefHelper::Create(
+          "my_func", {}, {}, {},
+          {{{"x"}, "UsesDefault", {}, {}},
+           {{"y"}, "ChangedFromDefault", {}, {{"a", 99}}}},
+          {});
+  TF_ASSERT_OK(NodeDefBuilder("call_func", "my_func", &function_registry)
+                   .Finalize(expected_graph_def.add_node()));
+  TF_EXPECT_GRAPH_EQ(expected_graph_def, produced_graph_def);
+  EXPECT_EQ(expected_graph_def.library().DebugString(),
+            produced_graph_def.library().DebugString());
+
+  std::set<std::pair<string, string>> expected_removed({{"UsesDefault", "a"}});
+  EXPECT_EQ(expected_removed, op_attr_removed);
+}
+
+TEST(StrippedOpListForGraphTest, FlatTest) {
+  // Make four ops
+  OpList op_list;
+  for (const string& op : {"A", "B", "C", "D"}) {
+    OpDef* op_def = op_list.add_op();
+    op_def->set_name(op);
+    op_def->set_summary("summary");
+    op_def->set_description("description");
+    op_def->set_is_commutative(op == "B");
+  }
+
+  // Make a graph which uses two ops once and twice, respectively.
+  // The result should be independent of the ordering.
+  const string graph_ops[4][3] = {
+      {"C", "B", "B"}, {"B", "C", "B"}, {"B", "B", "C"}, {"C", "C", "B"}};
+  for (const bool use_function : {false, true}) {
+    for (int order = 0; order < 4; order++) {
+      GraphDef graph_def;
+      if (use_function) {
+        FunctionDef* function_def = graph_def.mutable_library()->add_function();
+        function_def->mutable_signature()->set_name("F");
+        for (const string& op : graph_ops[order]) {
+          function_def->add_node_def()->set_op(op);
+        }
+        graph_def.add_node()->set_op("F");
+      } else {
+        for (const string& op : graph_ops[order]) {
+          string name = strings::StrCat("name", graph_def.node_size());
+          NodeDef* node = graph_def.add_node();
+          node->set_name(name);
+          node->set_op(op);
+        }
+      }
+
+      // Strip the op list
+      OpList stripped_op_list;
+      TF_ASSERT_OK(StrippedOpListForGraph(graph_def, OpListOpRegistry(&op_list),
+                                          &stripped_op_list));
+
+      // We should have exactly two ops: B and C.
+      ASSERT_EQ(stripped_op_list.op_size(), 2);
+      for (int i = 0; i < 2; i++) {
+        const OpDef& op = stripped_op_list.op(i);
+        EXPECT_EQ(op.name(), i ? "C" : "B");
+        EXPECT_EQ(op.summary(), "");
+        EXPECT_EQ(op.description(), "");
+        EXPECT_EQ(op.is_commutative(), !i);
+      }
+
+      // Should get the same result using OpsUsedByGraph().
+      std::set<string> used_ops;
+      OpsUsedByGraph(graph_def, &used_ops);
+      ASSERT_EQ(std::set<string>({"B", "C"}), used_ops);
+    }
+  }
+}
+
+TEST(StrippedOpListForGraphTest, NestedFunctionTest) {
+  // Make a primitive op A.
+  OpList op_list;
+  op_list.add_op()->set_name("A");
+
+  for (const bool recursive : {false, true}) {
+    // Call A from function B, and B from function C.
+    GraphDef graph_def;
+    FunctionDef* b = graph_def.mutable_library()->add_function();
+    FunctionDef* c = graph_def.mutable_library()->add_function();
+    b->mutable_signature()->set_name("B");
+    c->mutable_signature()->set_name("C");
+    b->add_node_def()->set_op("A");
+    c->add_node_def()->set_op("B");
+    if (recursive) {
+      b->add_node_def()->set_op("B");
+      c->add_node_def()->set_op("C");
+    }
+
+    // Use C in the graph.
+    graph_def.add_node()->set_op("C");
+
+    // The stripped op list should contain just A.
+    OpList stripped_op_list;
+    TF_ASSERT_OK(StrippedOpListForGraph(graph_def, OpListOpRegistry(&op_list),
+                                        &stripped_op_list));
+    ASSERT_EQ(stripped_op_list.op_size(), 1);
+    ASSERT_EQ(stripped_op_list.op(0).name(), "A");
+
+    // Should get the same result using OpsUsedByGraph().
+    std::set<string> used_ops;
+    OpsUsedByGraph(graph_def, &used_ops);
+    ASSERT_EQ(std::set<string>({"A"}), used_ops);
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow

graph_transfer_info.proto

@@ -0,0 +1,68 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "GraphTransferInfoProto";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/types.proto";
+
+// Protocol buffer representing a handle to a tensorflow resource. Handles are
+// not valid across executions, but can be serialized back and forth from within
+// a single run.
+message GraphTransferInfo {
+  enum Destination {
+    NOP = 0;
+    HEXAGON = 1;
+  }
+  message NodeInput {
+    int32 node_id = 1;
+    int32 output_port = 2;
+  }
+  message NodeInfo {
+    string name = 1;
+    int32 node_id = 2;
+    string type_name = 3;
+    int32 soc_op_id = 4;
+    int32 padding_id = 5;
+    int32 input_count = 6;
+    int32 output_count = 7;
+  };
+  message ConstNodeInfo {
+    string name = 1;
+    int32 node_id = 2;
+    repeated int64 shape = 3;
+    bytes data = 4;
+    DataType dtype = 5;
+  };
+  message NodeInputInfo {
+    int32 node_id = 1;
+    repeated NodeInput node_input = 2;
+  };
+  message NodeOutputInfo {
+    int32 node_id = 1;
+    repeated int32 max_byte_size = 2;
+  };
+  message GraphInputNodeInfo {
+    string name = 1;
+    repeated int64 shape = 2;
+    DataType dtype = 3;
+  }
+
+  message GraphOutputNodeInfo {
+    string name = 1;
+    repeated int64 shape = 2;
+    DataType dtype = 3;
+  }
+
+  repeated NodeInfo node_info = 1;
+  repeated ConstNodeInfo const_node_info = 2;
+  repeated NodeInputInfo node_input_info = 3;
+  repeated NodeOutputInfo node_output_info = 4;
+  // Input Node parameters of transferred graph
+  repeated GraphInputNodeInfo graph_input_node_info = 5;
+  repeated GraphOutputNodeInfo graph_output_node_info = 6;
+  // Destination of graph transfer
+  Destination destination = 7;
+};

iterator.proto

@@ -0,0 +1,17 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "IteratorProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.util";
+
+// Protocol buffer representing the metadata for an iterator's state stored
+// as a Variant tensor.
+message IteratorStateMetadata {
+  // A user-specified version string.
+  string version = 1;
+
+  // Keys for tensors in the VariantTensorDataProto.
+  repeated string keys = 2;
+}

kernel_def.proto

@@ -0,0 +1,36 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "KernelDefProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/attr_value.proto";
+
+message KernelDef {
+  // Must match the name of an Op.
+  string op = 1;
+
+  // Type of device this kernel runs on.
+  string device_type = 2;
+
+  message AttrConstraint {
+    // Name of an attr from the Op.
+    string name = 1;
+
+    // A list of values that this kernel supports for this attr.
+    // Like OpDef.AttrDef.allowed_values, except for kernels instead of Ops.
+    AttrValue allowed_values = 2;
+  }
+  repeated AttrConstraint constraint = 3;
+
+  // Names of the Op's input_/output_args that reside in host memory
+  // instead of device memory.
+  repeated string host_memory_arg = 4;
+
+  // This allows experimental kernels to be registered for an op that
+  // won't be used unless the user specifies a "_kernel" attr with
+  // value matching this.
+  string label = 5;
+}

kernel_def_builder.cc

@@ -0,0 +1,75 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/kernel_def.pb_text.h"
+#include "tensorflow/core/framework/kernel_def.pb.h"
+
+namespace tensorflow {
+
+KernelDefBuilder::KernelDefBuilder(const char* op_name) {
+  kernel_def_ = new KernelDef;
+  kernel_def_->set_op(op_name);
+}
+
+KernelDefBuilder::~KernelDefBuilder() {
+  DCHECK(kernel_def_ == nullptr) << "Did not call Build()";
+}
+
+KernelDefBuilder& KernelDefBuilder::Device(const char* device_type) {
+  kernel_def_->set_device_type(device_type);
+  return *this;
+}
+
+KernelDefBuilder& KernelDefBuilder::TypeConstraint(
+    const char* attr_name, gtl::ArraySlice<DataType> allowed) {
+  auto* constraint = kernel_def_->add_constraint();
+  constraint->set_name(attr_name);
+  auto* allowed_values = constraint->mutable_allowed_values()->mutable_list();
+  for (DataType dt : allowed) {
+    allowed_values->add_type(dt);
+  }
+  return *this;
+}
+
+KernelDefBuilder& KernelDefBuilder::TypeConstraint(const char* attr_name,
+                                                   DataType allowed) {
+  auto* constraint = kernel_def_->add_constraint();
+  constraint->set_name(attr_name);
+  constraint->mutable_allowed_values()->mutable_list()->add_type(allowed);
+  return *this;
+}
+
+KernelDefBuilder& KernelDefBuilder::HostMemory(const char* arg_name) {
+  kernel_def_->add_host_memory_arg(arg_name);
+  return *this;
+}
+
+KernelDefBuilder& KernelDefBuilder::Label(const char* label) {
+  CHECK_EQ(kernel_def_->label(), "")
+      << "Trying to set a kernel's label a second time: '" << label
+      << "' in: " << ProtoShortDebugString(*kernel_def_);
+  kernel_def_->set_label(label);
+  return *this;
+}
+
+const KernelDef* KernelDefBuilder::Build() {
+  KernelDef* r = kernel_def_;
+  kernel_def_ = nullptr;
+  return r;
+}
+
+}  // namespace tensorflow

kernel_def_builder.h

@@ -0,0 +1,87 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_
+#define TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_
+
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Forward declare proto so that kernels don't need to depend on it
+class KernelDef;
+
+// Builder class passed to the REGISTER_KERNEL_BUILDER() macro.
+class KernelDefBuilder {
+ public:
+  // Starts with just the name field set.
+  // Caller MUST call Build() and take ownership of the result.
+  explicit KernelDefBuilder(const char* op_name);
+  ~KernelDefBuilder();
+
+  // Required: specify the type of device this kernel supports.
+  // Returns *this.
+  KernelDefBuilder& Device(const char* device_type);
+  //  KernelDefBuilder& Device(DeviceType device_type);
+
+  // Specify that this kernel supports a limited set of values for a
+  // particular type or list(type) attr (a further restriction than
+  // what the Op allows).
+  // Returns *this.
+  KernelDefBuilder& TypeConstraint(const char* attr_name,
+                                   gtl::ArraySlice<DataType> allowed);
+
+  // Like TypeConstraint but supports just a single type.
+  KernelDefBuilder& TypeConstraint(const char* attr_name, DataType allowed);
+
+  // Like TypeConstraint, but (a) gets the type from a template parameter
+  // and (b) only supports a constraint to a single type.
+  template <class T>
+  KernelDefBuilder& TypeConstraint(const char* attr_name);
+  // TODO(josh11b): Support other types of attr constraints as needed.
+
+  // Specify that this kernel requires/provides an input/output arg
+  // in host memory (instead of the default, device memory).
+  // Returns *this.
+  KernelDefBuilder& HostMemory(const char* arg_name);
+
+  // Specify that this kernel requires a particular value for the
+  // "_kernel" attr.  May only be specified once.  Returns *this.
+  KernelDefBuilder& Label(const char* label);
+
+  // Returns a pointer to a KernelDef with fields set based on the
+  // above calls to this instance.
+  // Caller takes ownership of the result.
+  const KernelDef* Build();
+
+ private:
+  KernelDef* kernel_def_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(KernelDefBuilder);
+};
+
+// IMPLEMENTATION
+
+template <class T>
+KernelDefBuilder& KernelDefBuilder::TypeConstraint(const char* attr_name) {
+  return this->TypeConstraint(attr_name, DataTypeToEnum<T>::v());
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_KERNEL_DEF_BUILDER_H_

kernel_def_builder_test.cc

@@ -0,0 +1,91 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/kernel_def_builder.h"
+
+#include "tensorflow/core/framework/kernel_def.pb.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(KernelDefBuilderTest, Basic) {
+  const KernelDef* def = KernelDefBuilder("A").Device(DEVICE_CPU).Build();
+  KernelDef expected;
+  protobuf::TextFormat::ParseFromString("op: 'A' device_type: 'CPU'",
+                                        &expected);
+  EXPECT_EQ(def->DebugString(), expected.DebugString());
+  delete def;
+}
+
+TEST(KernelDefBuilderTest, TypeConstraint) {
+  const KernelDef* def = KernelDefBuilder("B")
+                             .Device(DEVICE_GPU)
+                             .TypeConstraint<float>("T")
+                             .Build();
+  KernelDef expected;
+  protobuf::TextFormat::ParseFromString(R"proto(
+    op: 'B' device_type: 'GPU'
+    constraint { name: 'T' allowed_values { list { type: DT_FLOAT } } } )proto",
+                                        &expected);
+
+  EXPECT_EQ(def->DebugString(), expected.DebugString());
+  delete def;
+
+  def = KernelDefBuilder("C")
+            .Device(DEVICE_GPU)
+            .TypeConstraint<int32>("U")
+            .TypeConstraint<bool>("V")
+            .Build();
+
+  protobuf::TextFormat::ParseFromString(R"proto(
+    op: 'C' device_type: 'GPU'
+    constraint { name: 'U' allowed_values { list { type: DT_INT32 } } }
+    constraint { name: 'V' allowed_values { list { type: DT_BOOL } } } )proto",
+                                        &expected);
+  EXPECT_EQ(def->DebugString(), expected.DebugString());
+  delete def;
+
+  def = KernelDefBuilder("D")
+            .Device(DEVICE_CPU)
+            .TypeConstraint("W", {DT_DOUBLE, DT_STRING})
+            .Build();
+  protobuf::TextFormat::ParseFromString(R"proto(
+    op: 'D' device_type: 'CPU'
+    constraint { name: 'W'
+        allowed_values { list { type: [DT_DOUBLE, DT_STRING] } } } )proto",
+                                        &expected);
+  EXPECT_EQ(def->DebugString(), expected.DebugString());
+  delete def;
+}
+
+TEST(KernelDefBuilderTest, HostMemory) {
+  const KernelDef* def = KernelDefBuilder("E")
+                             .Device(DEVICE_GPU)
+                             .HostMemory("in")
+                             .HostMemory("out")
+                             .Build();
+  KernelDef expected;
+  protobuf::TextFormat::ParseFromString(
+      "op: 'E' device_type: 'GPU' "
+      "host_memory_arg: ['in', 'out']",
+      &expected);
+  EXPECT_EQ(def->DebugString(), expected.DebugString());
+  delete def;
+}
+
+}  // namespace
+}  // namespace tensorflow

load_library.cc

@@ -0,0 +1,104 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include <memory>
+#include <unordered_set>
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mem.h"
+
+namespace tensorflow {
+
+namespace {
+
+struct Library {
+  void* handle = nullptr;
+  OpList op_list;
+};
+
+}  // namespace
+
+// Load a dynamic library.
+// On success, returns the handle to library in result, copies the serialized
+// OpList of OpDefs registered in the library to *buf and the length to *len,
+// and returns OK from the function. Otherwise return nullptr in result
+// and an error status from the function, leaving buf and len untouched.
+//
+// If `library_filename` has already been loaded, we return a cached handle
+// and OpList. Ops and kernels are registered as globals when a library is
+// loaded for the first time. Without caching, every subsequent load would not
+// perform initialization again, so the OpList would be empty.
+Status LoadLibrary(const char* library_filename, void** result,
+                   const void** buf, size_t* len) {
+  static mutex mu(LINKER_INITIALIZED);
+  static std::unordered_map<string, Library> loaded_libs;
+  Env* env = Env::Default();
+  Library library;
+  std::unordered_set<string> seen_op_names;
+  {
+    mutex_lock lock(mu);
+    if (loaded_libs.find(library_filename) != loaded_libs.end()) {
+      library = loaded_libs[library_filename];
+    } else {
+      Status s = OpRegistry::Global()->ProcessRegistrations();
+      if (!s.ok()) {
+        return s;
+      }
+      TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(
+          [&library, &seen_op_names](const Status& s,
+                                     const OpDef& opdef) -> Status {
+            if (errors::IsAlreadyExists(s)) {
+              if (seen_op_names.find(opdef.name()) == seen_op_names.end()) {
+                // Over writing a registration of an op not in this custom op
+                // library. Treat this as not an error.
+                return Status::OK();
+              }
+            }
+            if (s.ok()) {
+              *library.op_list.add_op() = opdef;
+              seen_op_names.insert(opdef.name());
+            }
+            return s;
+          }));
+      OpRegistry::Global()->DeferRegistrations();
+      s = env->LoadLibrary(library_filename, &library.handle);
+      if (s.ok()) {
+        s = OpRegistry::Global()->ProcessRegistrations();
+      }
+      if (!s.ok()) {
+        OpRegistry::Global()->ClearDeferredRegistrations();
+        TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(nullptr));
+        return s;
+      }
+      TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(nullptr));
+
+      loaded_libs[library_filename] = library;
+    }
+  }
+  string str;
+  library.op_list.SerializeToString(&str);
+  char* str_buf = reinterpret_cast<char*>(port::Malloc(str.length()));
+  memcpy(str_buf, str.data(), str.length());
+  *buf = str_buf;
+  *len = str.length();
+
+  *result = library.handle;
+  return Status::OK();
+}
+
+}  // namespace tensorflow

log_memory.cc

@@ -0,0 +1,102 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/log_memory.h"
+
+#include "tensorflow/core/framework/log_memory.pb_text.h"
+#include "tensorflow/core/framework/log_memory.pb.h"
+
+namespace tensorflow {
+
+const string LogMemory::kLogMemoryLabel = "__LOG_MEMORY__";
+
+bool LogMemory::IsEnabled() { return VLOG_IS_ON(1); }
+
+namespace {
+
+// Write the proto entry to LOG(INFO).
+template <typename T>
+void OutputToLog(const T& proto) {
+  string type_name = proto.GetTypeName();
+  const size_t index = type_name.find_last_of(".");
+  if (index != string::npos) type_name = type_name.substr(index + 1);
+  LOG(INFO) << LogMemory::kLogMemoryLabel << " " << type_name << " { "
+            << ProtoShortDebugString(proto) << " }";
+}
+
+}  // namespace
+
+void LogMemory::RecordStep(const int64 step_id, const string& handle) {
+  MemoryLogStep step;
+  step.set_step_id(step_id);
+  step.set_handle(handle);
+  OutputToLog(step);
+}
+
+void LogMemory::RecordTensorAllocation(const string& kernel_name,
+                                       const int64 step_id,
+                                       const Tensor& tensor) {
+  MemoryLogTensorAllocation allocation;
+  allocation.set_step_id(step_id);
+  allocation.set_kernel_name(kernel_name);
+  tensor.FillDescription(allocation.mutable_tensor());
+  OutputToLog(allocation);
+}
+
+void LogMemory::RecordTensorDeallocation(const int64 allocation_id,
+                                         const string& allocator_name) {
+  MemoryLogTensorDeallocation deallocation;
+  deallocation.set_allocation_id(allocation_id);
+  deallocation.set_allocator_name(allocator_name);
+  OutputToLog(deallocation);
+}
+
+void LogMemory::RecordTensorOutput(const string& kernel_name,
+                                   const int64 step_id, const int index,
+                                   const Tensor& tensor) {
+  MemoryLogTensorOutput output;
+  output.set_step_id(step_id);
+  output.set_kernel_name(kernel_name);
+  output.set_index(index);
+  tensor.FillDescription(output.mutable_tensor());
+  OutputToLog(output);
+}
+
+void LogMemory::RecordRawAllocation(const string& operation,
+                                    const int64 step_id, size_t num_bytes,
+                                    void* ptr, Allocator* allocator) {
+  MemoryLogRawAllocation allocation;
+  allocation.set_step_id(step_id);
+  allocation.set_operation(operation);
+  allocation.set_num_bytes(static_cast<int64>(num_bytes));
+  allocation.set_ptr(reinterpret_cast<uintptr_t>(ptr));
+  allocation.set_allocation_id(allocator->AllocationId(ptr));
+  allocation.set_allocator_name(allocator->Name());
+  OutputToLog(allocation);
+}
+
+void LogMemory::RecordRawDeallocation(const string& operation,
+                                      const int64 step_id, void* ptr,
+                                      Allocator* allocator, bool deferred) {
+  MemoryLogRawDeallocation deallocation;
+  deallocation.set_step_id(step_id);
+  deallocation.set_operation(operation);
+  deallocation.set_allocation_id(allocator->AllocationId(ptr));
+  deallocation.set_allocator_name(allocator->Name());
+  deallocation.set_deferred(deferred);
+  OutputToLog(deallocation);
+}
+
+}  // namespace tensorflow

log_memory.h

@@ -0,0 +1,111 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_LOG_MEMORY_H_
+#define TENSORFLOW_FRAMEWORK_LOG_MEMORY_H_
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+
+// LogMemory contains methods for recording memory allocations and
+// frees, associating each allocation with a step identified by a
+// process-wide id. For now, logging is enabled whenever VLOG_IS_ON(1)
+// for the log_memory module.
+//
+// Limitations: We don't log memory allocations by Eigen on the CPU
+// since that would require major changes to plumb through to the
+// Eigen::{DefaultDevice,ThreadPoolDevice} allocate and deallocate
+// methods. We do log Eigen allocations on GPU since the plumbing was
+// already in place.
+class LogMemory {
+ public:
+  // Allocations sometimes happen outside any computation step, and
+  // SpecialStepIds lists the ids used for those steps.
+  enum SpecialStepIds {
+    // Used when performing a just-in-time constant folding optimization.
+    CONSTANT_FOLDING_STEP_ID = -1,
+    // Used when constructing an Op kernel before executing a step.
+    OP_KERNEL_CONSTRUCTION_STEP_ID = -2,
+    // Used when allocating a tensor buffer from external code, e.g.,
+    // the C API.
+    EXTERNAL_TENSOR_ALLOCATION_STEP_ID = -3,
+    // Used when allocating a buffer for network transfer.
+    NETWORK_BUFFER_STEP_ID = -4,
+    // Used when allocating a buffer to fill a Proto from the GPU.
+    PROTO_BUFFER_STEP_ID = -5,
+    // Used when allocating a Tensor where the caller has not indicated
+    // the step.
+    UNKNOWN_STEP_ID = -6,
+  };
+
+  static const string kLogMemoryLabel;
+
+  // Test to see if memory logging is enabled. For now, logging is
+  // enabled whenever VLOG_IS_ON(1) for the log_memory module.
+  static bool IsEnabled();
+
+  // Log the beginning of a step.
+  static void RecordStep(int64 step_id, const string& handle);
+
+  // Log a tensor buffer allocation. The name indicates which kernel
+  // made the allocation. If the allocation is made through an
+  // OpKernelContext the step_id indicates which step is executing,
+  // otherwise step_id is one of the SpecialStepIds defined in
+  // op_kernel.h, e.g. Op Kernel construction or an optimization pass
+  // such as constant folding.
+  static void RecordTensorAllocation(const string& kernel_name, int64 step_id,
+                                     const Tensor& tensor);
+
+  // Log a tensor buffer deallocation. The deallocation is triggered
+  // when the buffer's refcount falls to zero, and the tracking
+  // mechanism does not associate it with a particular step or
+  // kernel. The allocation_id/allocator_name should match a
+  // corresponding tensor previously passed in to
+  // RecordTensorAllocation.
+  static void RecordTensorDeallocation(int64 allocation_id,
+                                       const string& allocator_name);
+
+  // Log the use of a tensor as an output from a kernel.
+  static void RecordTensorOutput(const string& kernel_name, int64 step_id,
+                                 int index, const Tensor& tensor);
+
+  // Log a "raw" allocation, which is just a buffer sized in
+  // bytes. The Eigen allocator, and memory copies, record their
+  // allocations this way, since they do not allocate TensorFlow
+  // tensors. The operation is set to the OpKernel name if this is
+  // called from within an Op execution, otherwise it indicates an
+  // operation such as memcpy. The step_id if >=0 indicates which step
+  // is executing, otherwise step_id is one of the SpecialStepIds
+  // defined in op_kernel.h, e.g. Op Kernel construction or an
+  // optimization pass such as constant folding.
+  static void RecordRawAllocation(const string& operation, int64 step_id,
+                                  size_t num_bytes, void* ptr,
+                                  Allocator* allocator);
+
+  // Log a "raw" deallocation of a buffer. When deferred is true, the
+  // buffer won't be used again, but a GPU kernel may still be
+  // enqueued using the buffer. A deferred deallocation should always
+  // be followed by a matching non-deferred deallocation when the
+  // buffer is actually returned and can be reused.
+  static void RecordRawDeallocation(const string& operation, int64 step_id,
+                                    void* ptr, Allocator* allocator,
+                                    bool deferred);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_LOG_MEMORY_H_

log_memory.proto

@@ -0,0 +1,93 @@
+syntax = "proto3";
+
+package tensorflow;
+option cc_enable_arenas = true;
+option java_outer_classname = "LogMemoryProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+
+import "tensorflow/core/framework/tensor_description.proto";
+
+message MemoryLogStep {
+  // Process-unique step id.
+  int64 step_id = 1;
+
+  // Handle describing the feeds and fetches of the step.
+  string handle = 2;
+};
+
+message MemoryLogTensorAllocation {
+  // Process-unique step id.
+  int64 step_id = 1;
+
+  // Name of the kernel making the allocation as set in GraphDef,
+  // e.g., "affine2/weights/Assign".
+  string kernel_name = 2;
+
+  // Allocated tensor details.
+  TensorDescription tensor = 3;
+};
+
+message MemoryLogTensorDeallocation {
+  // Id of the tensor buffer being deallocated, used to match to a
+  // corresponding allocation.
+  int64 allocation_id = 1;
+
+  // Name of the allocator used.
+  string allocator_name = 2;
+};
+
+message MemoryLogTensorOutput {
+  // Process-unique step id.
+  int64 step_id = 1;
+
+  // Name of the kernel producing an output as set in GraphDef, e.g.,
+  // "affine2/weights/Assign".
+  string kernel_name = 2;
+
+  // Index of the output being set.
+  int32 index = 3;
+
+  // Output tensor details.
+  TensorDescription tensor = 4;
+}
+
+message MemoryLogRawAllocation {
+  // Process-unique step id.
+  int64 step_id = 1;
+
+  // Name of the operation making the allocation.
+  string operation = 2;
+
+  // Number of bytes in the allocation.
+  int64 num_bytes = 3;
+
+  // Address of the allocation.
+  uint64 ptr = 4;
+
+  // Id of the tensor buffer being allocated, used to match to a
+  // corresponding deallocation.
+  int64 allocation_id = 5;
+
+  // Name of the allocator used.
+  string allocator_name = 6;
+};
+
+message MemoryLogRawDeallocation {
+  // Process-unique step id.
+  int64 step_id = 1;
+
+  // Name of the operation making the deallocation.
+  string operation = 2;
+
+  // Id of the tensor buffer being deallocated, used to match to a
+  // corresponding allocation.
+  int64 allocation_id = 3;
+
+  // Name of the allocator used.
+  string allocator_name = 4;
+
+  // True if the deallocation is queued and will be performed later,
+  // e.g. for GPU lazy freeing of buffers.
+  bool deferred = 5;
+};

lookup_interface.cc

@@ -0,0 +1,87 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/lookup_interface.h"
+
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace lookup {
+
+Status LookupInterface::CheckKeyShape(const TensorShape& shape) {
+  if (!TensorShapeUtils::EndsWith(shape, key_shape())) {
+    return errors::InvalidArgument("Input key shape ", shape.DebugString(),
+                                   " must end with the table's key shape ",
+                                   key_shape().DebugString());
+  }
+  return Status::OK();
+}
+
+Status LookupInterface::CheckKeyAndValueTypes(const Tensor& keys,
+                                              const Tensor& values) {
+  if (keys.dtype() != key_dtype()) {
+    return errors::InvalidArgument("Key must be type ", key_dtype(),
+                                   " but got ", keys.dtype());
+  }
+  if (values.dtype() != value_dtype()) {
+    return errors::InvalidArgument("Value must be type ", value_dtype(),
+                                   " but got ", values.dtype());
+  }
+  return Status::OK();
+}
+
+Status LookupInterface::CheckKeyAndValueTensorsHelper(const Tensor& keys,
+                                                      const Tensor& values) {
+  TF_RETURN_IF_ERROR(CheckKeyAndValueTypes(keys, values));
+  TF_RETURN_IF_ERROR(CheckKeyShape(keys.shape()));
+
+  TensorShape expected_value_shape = keys.shape();
+  for (int i = 0; i < key_shape().dims(); ++i) {
+    expected_value_shape.RemoveDim(expected_value_shape.dims() - 1);
+  }
+  expected_value_shape.AppendShape(value_shape());
+  if (values.shape() != expected_value_shape) {
+    return errors::InvalidArgument(
+        "Expected shape ", expected_value_shape.DebugString(),
+        " for value, got ", values.shape().DebugString());
+  }
+  return Status::OK();
+}
+
+Status LookupInterface::CheckKeyAndValueTensorsForInsert(const Tensor& keys,
+                                                         const Tensor& values) {
+  return CheckKeyAndValueTensorsHelper(keys, values);
+}
+
+Status LookupInterface::CheckKeyAndValueTensorsForImport(const Tensor& keys,
+                                                         const Tensor& values) {
+  return CheckKeyAndValueTensorsHelper(keys, values);
+}
+
+Status LookupInterface::CheckFindArguments(const Tensor& key,
+                                           const Tensor& default_value) {
+  TF_RETURN_IF_ERROR(CheckKeyAndValueTypes(key, default_value));
+  TF_RETURN_IF_ERROR(CheckKeyShape(key.shape()));
+  if (default_value.shape() != value_shape()) {
+    return errors::InvalidArgument(
+        "Expected shape ", value_shape().DebugString(),
+        " for default value, got ", default_value.shape().DebugString());
+  }
+  return Status::OK();
+}
+
+}  // namespace lookup
+}  // namespace tensorflow

lookup_interface.h

@@ -0,0 +1,145 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_FRAMEWORK_LOOKUP_INTERFACE_H_
+#define TENSORFLOW_FRAMEWORK_LOOKUP_INTERFACE_H_
+
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+class OpKernelContext;
+
+namespace lookup {
+
+// Forward declaration so we can define GetInitializableLookupTable() in
+// LookupInterface.
+class InitializableLookupTable;
+
+// Lookup interface for batch lookups used by table lookup ops.
+class LookupInterface : public ResourceBase {
+ public:
+  // Performs batch lookups, for every element in the key tensor, Find returns
+  // the corresponding value into the values tensor.
+  // If an element is not present in the table, the given default value is used.
+
+  // For tables that require initialization, Find is available once the table
+  // is marked as initialized.
+
+  // Returns the following statuses:
+  // - OK: when the find finishes successfully.
+  // - FailedPrecondition: if the table is not initialized.
+  // - InvalidArgument: if any of the preconditions on the lookup key or value
+  //   fails.
+  // - In addition, other implementations may provide another non-OK status
+  //   specific to their failure modes.
+  virtual Status Find(OpKernelContext* ctx, const Tensor& keys, Tensor* values,
+                      const Tensor& default_value) = 0;
+
+  // Inserts elements into the table. Each element of the key tensor is
+  // associated with the corresponding element in the value tensor.
+  // This method is only implemented in mutable tables that can be updated over
+  // the execution of the graph. It returns Status::NotImplemented for read-only
+  // tables that are initialized once before they can be looked up.
+
+  // Returns the following statuses:
+  // - OK: when the insert finishes successfully.
+  // - InvalidArgument: if any of the preconditions on the lookup key or value
+  //   fails.
+  // - Unimplemented: if the table does not support insertions.
+  virtual Status Insert(OpKernelContext* ctx, const Tensor& keys,
+                        const Tensor& values) = 0;
+
+  // Returns the number of elements in the table.
+  virtual size_t size() const = 0;
+
+  // Exports the values of the table to two tensors named keys and values.
+  // Note that the shape of the tensors is completely up to the implementation
+  // of the table and can be different than the tensors used for the Insert
+  // function above.
+  virtual Status ExportValues(OpKernelContext* ctx) = 0;
+
+  // Imports previously exported keys and values.
+  // As mentioned above, the shape of the keys and values tensors are determined
+  // by the ExportValues function above and can be different than for the
+  // Insert function.
+  virtual Status ImportValues(OpKernelContext* ctx, const Tensor& keys,
+                              const Tensor& values) = 0;
+
+  // Returns the data type of the key.
+  virtual DataType key_dtype() const = 0;
+
+  // Returns the data type of the value.
+  virtual DataType value_dtype() const = 0;
+
+  // Returns the shape of a key in the table.
+  virtual TensorShape key_shape() const = 0;
+
+  // Returns the shape of a value in the table.
+  virtual TensorShape value_shape() const = 0;
+
+  // Check format of the key and value tensors for the Insert function.
+  // Returns OK if all the following requirements are satisfied, otherwise it
+  // returns InvalidArgument:
+  // - DataType of the tensor keys equals to the table key_dtype
+  // - DataType of the tensor values equals to the table value_dtype
+  // - the values tensor has the required shape given keys and the tables's
+  //   value shape.
+  virtual Status CheckKeyAndValueTensorsForInsert(const Tensor& keys,
+                                                  const Tensor& values);
+
+  // Similar to the function above but instead checks eligibility for the Import
+  // function.
+  virtual Status CheckKeyAndValueTensorsForImport(const Tensor& keys,
+                                                  const Tensor& values);
+
+  // Check the arguments of a find operation. Returns OK if all the following
+  // requirements are satisfied, otherwise it returns InvalidArgument:
+  // - DataType of the tensor keys equals to the table key_dtype
+  // - DataType of the tensor default_value equals to the table value_dtype
+  // - the default_value tensor shape matches the table's value shape.
+  Status CheckFindArguments(const Tensor& keys, const Tensor& default_value);
+
+  string DebugString() override {
+    return strings::StrCat("A lookup table of size: ", size());
+  }
+
+  // Returns an InitializableLookupTable, a subclass of LookupInterface, if the
+  // current object is an InitializableLookupTable. Otherwise, returns nullptr.
+  virtual InitializableLookupTable* GetInitializableLookupTable() {
+    return nullptr;
+  }
+
+ protected:
+  virtual ~LookupInterface() = default;
+
+  // Makes sure that the key and value tensor DataType's match the table
+  // key_dtype and value_dtype.
+  Status CheckKeyAndValueTypes(const Tensor& keys, const Tensor& values);
+
+  // Makes sure that the provided shape is consistent with the table keys shape.
+  Status CheckKeyShape(const TensorShape& shape);
+
+ private:
+  Status CheckKeyAndValueTensorsHelper(const Tensor& keys,
+                                       const Tensor& values);
+};
+
+}  // namespace lookup
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_FRAMEWORK_LOOKUP_INTERFACE_H_

memory_types.cc

@@ -0,0 +1,156 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/memory_types.h"
+
+#include <utility>
+
+#include "tensorflow/core/framework/kernel_def.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+namespace {
+// Returns the largest endpoint of anything in the name_map.
+int GetTotal(const NameRangeMap& name_map) {
+  int total = 0;
+  for (const auto& item : name_map) {
+    total = std::max(total, item.second.second);
+  }
+  return total;
+}
+
+// Fills memory_types for either input or output, setting everything
+// to DEVICE_MEMORY except those args in host_memory_args.  Removes
+// elements of host_memory_args that were used.
+void MemoryTypesHelper(const NameRangeMap& name_map,
+                       std::vector<string>* host_memory_args,
+                       MemoryTypeVector* memory_types) {
+  // Update args that have been marked as in "HOST_MEMORY".
+  size_t keep = 0;
+  for (size_t i = 0; i < host_memory_args->size(); ++i) {
+    auto iter = name_map.find((*host_memory_args)[i]);
+    if (iter != name_map.end()) {
+      for (int j = iter->second.first; j < iter->second.second; ++j) {
+        (*memory_types)[j] = HOST_MEMORY;
+      }
+    } else {
+      // (*host_memory_args)[i] not found, save it for the next pass.
+      if (i > keep) (*host_memory_args)[keep] = (*host_memory_args)[i];
+      ++keep;
+    }
+  }
+  host_memory_args->resize(keep);
+}
+
+MemoryType MTypeFromDType(const DataType dtype) {
+  return (dtype == DT_INT32 || DataTypeAlwaysOnHost(dtype)) ? HOST_MEMORY
+                                                            : DEVICE_MEMORY;
+}
+
+}  // namespace
+
+Status MemoryTypesForNode(const OpRegistryInterface* op_registry,
+                          const DeviceType& device_type, const NodeDef& ndef,
+                          MemoryTypeVector* inp_mtypes,
+                          MemoryTypeVector* out_mtypes) {
+  // Look up the Op registered for this op name.
+  const OpDef* op_def;
+  TF_RETURN_IF_ERROR(op_registry->LookUpOpDef(ndef.op(), &op_def));
+
+  // Look up the Kernel registered for this node def.
+  const KernelDef* kdef = nullptr;
+  Status status =
+      FindKernelDef(device_type, ndef, &kdef, nullptr /* kernel_class_name */);
+
+  DataTypeVector inp_dtypes;
+  DataTypeVector out_dtypes;
+  TF_RETURN_IF_ERROR(
+      InOutTypesForNode(ndef, *op_def, &inp_dtypes, &out_dtypes));
+
+  inp_mtypes->clear();
+  out_mtypes->clear();
+
+  // For functions (which have no KernelDef) and their gradients, we can only
+  // best-effort derive the memory type from the data type. For now, we assume
+  // int32 is always on host memory and other types are always on device memory.
+  // TODO(zhifengc,phawkins): We should do type inference over function bodies
+  // to derive the correct input/output memory types. We should also split
+  // host-memory and non host-memory arguments into separate type lists.
+  if (!status.ok() || ndef.op() == "SymbolicGradient") {
+    for (const auto& t : inp_dtypes) inp_mtypes->push_back(MTypeFromDType(t));
+    for (const auto& t : out_dtypes) out_mtypes->push_back(MTypeFromDType(t));
+    return Status::OK();
+  }
+
+  // Gets the input/output names and their corresponding endpoint ranges.
+  NameRangeMap inp_names;
+  NameRangeMap out_names;
+  TF_RETURN_IF_ERROR(NameRangesForNode(ndef, *op_def, &inp_names, &out_names));
+
+  // Now that we know the size, fill with the default 'DEVICE_MEMORY'.
+  inp_mtypes->resize(GetTotal(inp_names), DEVICE_MEMORY);
+  out_mtypes->resize(GetTotal(out_names), DEVICE_MEMORY);
+
+  // Fills in host memory types based on the kernel def.
+  const auto& from_proto = kdef->host_memory_arg();
+  std::vector<string> host_memory_args(from_proto.begin(), from_proto.end());
+  MemoryTypesHelper(inp_names, &host_memory_args, inp_mtypes);
+  MemoryTypesHelper(out_names, &host_memory_args, out_mtypes);
+  if (!host_memory_args.empty()) {
+    return errors::InvalidArgument(
+        "HostMemory args '", str_util::Join(host_memory_args, "', '"),
+        "' not found in OpDef: ", SummarizeOpDef(*op_def));
+  }
+  CHECK_LE(inp_mtypes->size(), inp_dtypes.size());
+  CHECK_LE(out_mtypes->size(), out_dtypes.size());
+
+  // Mark e.g. all resource and string types as host memory.
+  for (int i = 0; i < inp_mtypes->size(); ++i) {
+    if (DataTypeAlwaysOnHost(inp_dtypes[i])) {
+      (*inp_mtypes)[i] = HOST_MEMORY;
+    }
+  }
+  for (int i = 0; i < out_mtypes->size(); ++i) {
+    if (DataTypeAlwaysOnHost(out_dtypes[i])) {
+      (*out_mtypes)[i] = HOST_MEMORY;
+    }
+  }
+
+  std::vector<int32> hostmem_attr;
+  if (GetNodeAttr(ndef, "_input_hostmem", &hostmem_attr).ok()) {
+    for (int32 i : hostmem_attr) {
+      if (0 <= i && i < inp_mtypes->size()) {
+        (*inp_mtypes)[i] = HOST_MEMORY;
+      }
+    }
+  }
+  if (GetNodeAttr(ndef, "_output_hostmem", &hostmem_attr).ok()) {
+    for (int32 i : hostmem_attr) {
+      if (0 <= i && i < out_mtypes->size()) {
+        (*out_mtypes)[i] = HOST_MEMORY;
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace tensorflow