Export @195097388.

research/syntaxnet/dragnn/runtime/mmap.h

+// Copyright 2017 Google Inc. 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.
+// =============================================================================
+
+// Utils for establishing and managing memory-mapped files.
+
+#ifndef DRAGNN_RUNTIME_MMAP_H_
+#define DRAGNN_RUNTIME_MMAP_H_
+
+#include <stddef.h>
+#include <memory>
+#include <string>
+
+#include "dragnn/runtime/alignment.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// A uniquely-owned aligned memory-mapped file.  This has virtual methods only
+// for mocking in tests; do not derive from class.
+class UniqueAlignedMmap {
+ public:
+  // A mockable wrapper around the system calls used by this class.
+  class Syscalls {
+   public:
+    virtual ~Syscalls() = default;
+
+    // Each method below forwards to the similarly-named syscall.  Some methods
+    // have been simplified by omitting arguments that are never varied.
+    virtual int Open(const string &path) const;
+    virtual int Close(int file_descriptor) const;
+    virtual void *Mmap(int file_descriptor, size_t size) const;
+    virtual int Munmap(void *data, size_t size) const;
+  };
+
+  // Creates an empty, unmapped memory region.
+  UniqueAlignedMmap() = default;
+
+  // FOR TESTS ONLY.  As above, but injects the |syscalls|.
+  explicit UniqueAlignedMmap(std::unique_ptr<Syscalls> syscalls);
+
+  // Supports movement only.
+  UniqueAlignedMmap(UniqueAlignedMmap &&that);
+  UniqueAlignedMmap &operator=(UniqueAlignedMmap &&that);
+  UniqueAlignedMmap(const UniqueAlignedMmap &that) = delete;
+  UniqueAlignedMmap &operator=(const UniqueAlignedMmap &that) = delete;
+
+  // Unmaps the current memory-mapped file, if any.
+  ~UniqueAlignedMmap();
+
+  // Resets this to a memory-mapping of the |path|.  On error, returns non-OK
+  // and modifies nothing.
+  tensorflow::Status Reset(const string &path);
+
+  // Returns the mapped memory region.
+  AlignedView view() const { return AlignedView(view_); }
+
+ private:
+  // Unmaps [|data|,|data|+|size|), if non-empty.  Uses the |path| for error
+  // logging.  Does not return a status because none of the call sites could
+  // pass it along; they'd log it anyways.
+  void UnmapIfNonEmpty(void *data, size_t size, const string &path) const;
+
+  // The system calls used to perform the memory-mapping.
+  std::unique_ptr<Syscalls> syscalls_{new Syscalls()};
+
+  // The current memory-mapped file, or empty if unmapped.  Mutable to satisfy
+  // munmap(), which requires a non-const pointer---contents are not modified.
+  MutableAlignedView view_;
+
+  // The path to the current memory-mapped file, if any, for debug logging.
+  string path_;
+};
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MMAP_H_

research/syntaxnet/dragnn/runtime/mmap_array_variable_store.cc

+// Copyright 2017 Google Inc. 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 "dragnn/runtime/mmap_array_variable_store.h"
+
+#include <utility>
+
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+tensorflow::Status MmapArrayVariableStore::Reset(
+    const ArrayVariableStoreSpec &spec, const string &path) {
+  UniqueAlignedMmap data;
+  TF_RETURN_IF_ERROR(data.Reset(path));
+  TF_RETURN_IF_ERROR(ArrayVariableStore::Reset(spec, data.view()));
+
+  // Success; make modifications.
+  data_ = std::move(data);
+  return tensorflow::Status::OK();
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/mmap_array_variable_store.h

+// Copyright 2017 Google Inc. 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 DRAGNN_RUNTIME_MMAP_ARRAY_VARIABLE_STORE_H_
+#define DRAGNN_RUNTIME_MMAP_ARRAY_VARIABLE_STORE_H_
+
+#include <string>
+
+#include "dragnn/protos/runtime.pb.h"
+#include "dragnn/runtime/array_variable_store.h"
+#include "dragnn/runtime/mmap.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// An ArrayVariableStore subclass that maps file content into memory.
+class MmapArrayVariableStore : public ArrayVariableStore {
+ public:
+  // Creates an uninitialized store.
+  MmapArrayVariableStore() = default;
+
+  // Resets this to represent the variables defined by the |spec|, mapping the
+  // byte array from the |path|.  On error, returns non-OK and modifies nothing.
+  tensorflow::Status Reset(const ArrayVariableStoreSpec &spec,
+                           const string &path);
+
+ private:
+  // The memory-mapped file containing the variables.
+  UniqueAlignedMmap data_;
+};
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MMAP_ARRAY_VARIABLE_STORE_H_

research/syntaxnet/dragnn/runtime/mmap_test.cc

+// Copyright 2017 Google Inc. 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 "dragnn/runtime/mmap.h"
+
+#include <stddef.h>
+#include <sys/mman.h>
+#include <string>
+#include <utility>
+
+#include "dragnn/core/test/generic.h"
+#include "syntaxnet/base.h"
+#include <gmock/gmock.h>
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+using ::testing::Return;
+
+// A mockable set of system calls.
+class MockSyscalls : public UniqueAlignedMmap::Syscalls {
+ public:
+  MOCK_CONST_METHOD1(Open, int(const string &path));
+  MOCK_CONST_METHOD1(Close, int(int file_descriptor));
+  MOCK_CONST_METHOD2(Mmap, void *(int file_descriptor, size_t size));
+  MOCK_CONST_METHOD2(Munmap, int(void *, size_t size));
+};
+
+class UniqueAlignedMmapTest : public ::testing::Test {
+ protected:
+  const string kInvalidFile = "/some/invalid/path";
+  const string kEmptyFile = tensorflow::io::JoinPath(
+      test::GetTestDataPrefix(), "dragnn/runtime/testdata/empty_file");
+  const string kTenBytes = tensorflow::io::JoinPath(
+      test::GetTestDataPrefix(), "dragnn/runtime/testdata/ten_bytes");
+
+  std::unique_ptr<MockSyscalls> syscalls_{new MockSyscalls()};
+};
+
+// Tests that the mapped region is empty by default.
+TEST_F(UniqueAlignedMmapTest, EmptyByDefault) {
+  UniqueAlignedMmap data;
+  EXPECT_TRUE(data.view().empty());
+}
+
+// Tests that an empty file can be mapped.
+TEST_F(UniqueAlignedMmapTest, EmptyFile) {
+  UniqueAlignedMmap data;
+  TF_ASSERT_OK(data.Reset(kEmptyFile));
+  EXPECT_TRUE(data.view().empty());
+}
+
+// Tests that a non-empty file can be mapped.
+TEST_F(UniqueAlignedMmapTest, TenBytes) {
+  UniqueAlignedMmap data;
+  TF_ASSERT_OK(data.Reset(kTenBytes));
+  ASSERT_EQ(data.view().size(), 10);
+  EXPECT_STREQ(data.view().data(), "0123456789");
+}
+
+// Tests that the mapped files can be move-constructed and move-assigned.
+TEST_F(UniqueAlignedMmapTest, Movement) {
+  UniqueAlignedMmap data1;
+  TF_ASSERT_OK(data1.Reset(kTenBytes));
+
+  UniqueAlignedMmap data2(std::move(data1));
+  ASSERT_EQ(data2.view().size(), 10);
+  EXPECT_STREQ(data2.view().data(), "0123456789");
+
+  UniqueAlignedMmap data3;
+  data3 = std::move(data2);
+  ASSERT_EQ(data3.view().size(), 10);
+  EXPECT_STREQ(data3.view().data(), "0123456789");
+}
+
+// Tests that the mapping fails if the file is invalid.
+TEST_F(UniqueAlignedMmapTest, InvalidFile) {
+  UniqueAlignedMmap data;
+  EXPECT_FALSE(data.Reset(kInvalidFile).ok());
+}
+
+// Tests that the mapping fails if the file cannot be open()ed.
+TEST_F(UniqueAlignedMmapTest, FailToOpen) {
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(-1));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes), test::IsErrorWithSubstr("Failed to open"));
+}
+
+// Tests that the mapping fails if the file cannot be mmap()ed.
+TEST_F(UniqueAlignedMmapTest, FailToMmap) {
+  const int kFileDescriptor = 5;
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(kFileDescriptor));
+  EXPECT_CALL(*syscalls_, Mmap(kFileDescriptor, 10))
+      .WillOnce(Return(MAP_FAILED));
+  EXPECT_CALL(*syscalls_, Close(kFileDescriptor)).WillOnce(Return(0));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes), test::IsErrorWithSubstr("Failed to mmap"));
+}
+
+// As above, but also fails to close.
+TEST_F(UniqueAlignedMmapTest, FailToMmapAndClose) {
+  const int kFileDescriptor = 5;
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(kFileDescriptor));
+  EXPECT_CALL(*syscalls_, Mmap(kFileDescriptor, 10))
+      .WillOnce(Return(MAP_FAILED));
+  EXPECT_CALL(*syscalls_, Close(kFileDescriptor)).WillOnce(Return(-1));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes), test::IsErrorWithSubstr("Failed to mmap"));
+}
+
+// Tests that the mapping fails if the file cannot be close()ed.
+TEST_F(UniqueAlignedMmapTest, FailToClose) {
+  const int kFileDescriptor = 5;
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(kFileDescriptor));
+  EXPECT_CALL(*syscalls_, Mmap(kFileDescriptor, 10)).WillOnce(Return(nullptr));
+  EXPECT_CALL(*syscalls_, Close(kFileDescriptor)).WillOnce(Return(-1));
+  EXPECT_CALL(*syscalls_, Munmap(nullptr, 10)).WillOnce(Return(0));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes),
+              test::IsErrorWithSubstr("Failed to close"));
+}
+
+// As above, but also fails to munmap().
+TEST_F(UniqueAlignedMmapTest, FailToCloseAndMunmap) {
+  const int kFileDescriptor = 5;
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(kFileDescriptor));
+  EXPECT_CALL(*syscalls_, Mmap(kFileDescriptor, 10)).WillOnce(Return(nullptr));
+  EXPECT_CALL(*syscalls_, Close(kFileDescriptor)).WillOnce(Return(-1));
+  EXPECT_CALL(*syscalls_, Munmap(nullptr, 10)).WillOnce(Return(-1));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes),
+              test::IsErrorWithSubstr("Failed to close"));
+}
+
+// Tests that the mapping fails if the mapped region is misaligned.
+TEST_F(UniqueAlignedMmapTest, Misaligned) {
+  char *ptr = nullptr;
+  ++ptr;
+  const int kFileDescriptor = 5;
+  EXPECT_CALL(*syscalls_, Open(kTenBytes)).WillOnce(Return(kFileDescriptor));
+  EXPECT_CALL(*syscalls_, Mmap(kFileDescriptor, 10)).WillOnce(Return(ptr));
+  EXPECT_CALL(*syscalls_, Close(kFileDescriptor)).WillOnce(Return(0));
+  EXPECT_CALL(*syscalls_, Munmap(ptr, 10)).WillOnce(Return(0));
+
+  UniqueAlignedMmap data(std::move(syscalls_));
+  EXPECT_THAT(data.Reset(kTenBytes),
+              test::IsErrorWithSubstr("Pointer fails alignment requirement"));
+}
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/mst_solver_component_base.cc

+// Copyright 2018 Google Inc. 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 "dragnn/runtime/mst_solver_component_base.h"
+
+#include <stddef.h>
+
+#include "dragnn/runtime/attributes.h"
+#include "dragnn/runtime/math/types.h"
+#include "dragnn/runtime/network_unit.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+// Attributes used by the MST solver.
+struct MstSolverAttributes : public Attributes {
+  // Whether to solve for a spanning forest instead of a spanning tree.
+  Optional<bool> forest{"forest", false, this};
+
+  // Training-only attributes, ignored in the runtime.
+  Ignored loss{"loss", this};
+};
+
+}  // namespace
+
+MstSolverComponentBase::MstSolverComponentBase(const string &builder_name,
+                                               const string &backend_name)
+    : builder_name_(builder_name), backend_name_(backend_name) {}
+
+bool MstSolverComponentBase::Supports(
+    const ComponentSpec &component_spec,
+    const string &normalized_builder_name) const {
+  const string network_unit = NetworkUnit::GetClassName(component_spec);
+  return (normalized_builder_name == "BulkAnnotatorComponent" ||
+          normalized_builder_name == builder_name_) &&
+         (component_spec.backend().registered_name() == "StatelessComponent" ||
+          component_spec.backend().registered_name() == backend_name_) &&
+         component_spec.transition_system().registered_name() == "heads" &&
+         network_unit == "MstSolverNetwork" &&
+         component_spec.fixed_feature_size() == 0 &&
+         component_spec.linked_feature_size() == 1;
+}
+
+tensorflow::Status MstSolverComponentBase::Initialize(
+    const ComponentSpec &component_spec, VariableStore *variable_store,
+    NetworkStateManager *network_state_manager,
+    ExtensionManager *extension_manager) {
+  MstSolverAttributes attributes;
+  TF_RETURN_IF_ERROR(
+      attributes.Reset(component_spec.network_unit().parameters()));
+  forest_ = attributes.forest();
+
+  const LinkedFeatureChannel &link = component_spec.linked_feature(0);
+  size_t dimension = 0;
+  TF_RETURN_IF_ERROR(network_state_manager->LookupLayer(
+      link.source_component(), link.source_layer(), &dimension,
+      &adjacency_handle_));
+
+  if (dimension != 1) {
+    return tensorflow::errors::InvalidArgument(
+        "Adjacency matrix has dimension ", dimension, " but expected 1");
+  }
+
+  extension_manager->GetShared(&heads_handle_);
+  extension_manager->GetShared(&solver_handle_);
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MstSolverComponentBase::ComputeHeads(
+    SessionState *session_state,
+    tensorflow::gtl::ArraySlice<Index> *heads) const {
+  Matrix<float> adjacency(
+      session_state->network_states.GetLayer(adjacency_handle_));
+  const size_t num_nodes = adjacency.num_rows();
+
+  Solver &solver = session_state->extensions.Get(solver_handle_);
+  TF_RETURN_IF_ERROR(solver.Init(forest_, num_nodes));
+
+  for (size_t target = 0; target < num_nodes; ++target) {
+    Vector<float> source_scores = adjacency.row(target);
+    for (size_t source = 0; source < num_nodes; ++source) {
+      if (source == target) {
+        solver.AddRoot(source, source_scores[source]);
+      } else {
+        solver.AddArc(source, target, source_scores[source]);
+      }
+    }
+  }
+
+  std::vector<Index> &argmax = session_state->extensions.Get(heads_handle_);
+  argmax.resize(num_nodes);
+  TF_RETURN_IF_ERROR(solver.Solve(&argmax));
+
+  *heads = argmax;
+  return tensorflow::Status::OK();
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/mst_solver_component_base.h

+// Copyright 2018 Google Inc. 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 DRAGNN_RUNTIME_MST_SOLVER_COMPONENT_BASE_H_
+#define DRAGNN_RUNTIME_MST_SOLVER_COMPONENT_BASE_H_
+
+#include <vector>
+
+#include "dragnn/core/compute_session.h"
+#include "dragnn/mst/mst_solver.h"
+#include "dragnn/protos/spec.pb.h"
+#include "dragnn/runtime/component.h"
+#include "dragnn/runtime/extensions.h"
+#include "dragnn/runtime/network_states.h"
+#include "dragnn/runtime/session_state.h"
+#include "dragnn/runtime/variable_store.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// Base class for MST parsing components, which select heads jointly by finding
+// the maximum spanning tree of the input tokens.
+//
+// This base class only computes the selected heads, while subclasses apply the
+// heads to the annotations in the ComputeSession.
+class MstSolverComponentBase : public Component {
+ public:
+  // NB: This definition of Index should match the MstSolver TF op wrappers.
+  using Index = uint16;
+
+  // Partially implements Component.
+  bool Supports(const ComponentSpec &component_spec,
+                const string &normalized_builder_name) const override;
+  tensorflow::Status Initialize(const ComponentSpec &component_spec,
+                                VariableStore *variable_store,
+                                NetworkStateManager *network_state_manager,
+                                ExtensionManager *extension_manager) override;
+  bool PreferredTo(const Component &other) const override { return false; }
+
+ protected:
+  // Creates a component that supports the |builder_name| and |backend_name|.
+  MstSolverComponentBase(const string &builder_name,
+                         const string &backend_name);
+
+  // Points |heads| at the list of heads computed from the |session_state|,
+  // where a self-loop indicates a root.  Returns non-OK on error.
+  tensorflow::Status ComputeHeads(
+      SessionState *session_state,
+      tensorflow::gtl::ArraySlice<Index> *heads) const;
+
+ private:
+  using Solver = MstSolver<Index, float>;
+
+  // Names of the supported component builder and backend.
+  const string builder_name_;
+  const string backend_name_;
+
+  // Whether to solve for a spanning forest instead of a spanning tree.
+  bool forest_ = false;
+
+  // Directed adjacency matrix input.
+  PairwiseLayerHandle<float> adjacency_handle_;
+
+  // List of selected head indices.
+  SharedExtensionHandle<std::vector<Index>> heads_handle_;
+
+  // Reusable MST solver.
+  SharedExtensionHandle<Solver> solver_handle_;
+};
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MST_SOLVER_COMPONENT_BASE_H_

research/syntaxnet/dragnn/runtime/mst_solver_component_base_test.cc

+// Copyright 2018 Google Inc. 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 "dragnn/runtime/mst_solver_component_base.h"
+
+#include <stddef.h>
+#include <string>
+#include <vector>
+
+#include "dragnn/core/compute_session.h"
+#include "dragnn/core/input_batch_cache.h"
+#include "dragnn/core/test/generic.h"
+#include "dragnn/protos/spec.pb.h"
+#include "dragnn/protos/trace.pb.h"
+#include "dragnn/runtime/network_states.h"
+#include "dragnn/runtime/session_state.h"
+#include "dragnn/runtime/test/network_test_base.h"
+#include "dragnn/runtime/variable_store.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+constexpr size_t kNumSteps = 12;
+constexpr size_t kRootIndex = 7;  // the root and head of all other tokens
+
+constexpr char kTestBuilder[] = "TestBuilder";
+constexpr char kTestBackend[] = "TestBackend";
+constexpr char kPreviousComponentName[] = "previous_component";
+constexpr char kAdjacencyLayerName[] = "adjacency_layer";
+constexpr char kBadDimLayerName[] = "bad_layer";
+
+// A subclass for tests.
+class BasicMstSolverComponent : public MstSolverComponentBase {
+ public:
+  BasicMstSolverComponent()
+      : MstSolverComponentBase(kTestBuilder, kTestBackend) {}
+
+  // Implements Component.  These methods are never called, but must be defined
+  // so the class is not abstract.
+  tensorflow::Status Evaluate(
+      SessionState *session_state, ComputeSession *compute_session,
+      ComponentTrace *component_trace) const override {
+    return tensorflow::Status::OK();
+  }
+
+  // Publicizes the base class's method.
+  using MstSolverComponentBase::ComputeHeads;
+};
+
+// Returns a ComponentSpec that works with the head selection component.
+ComponentSpec MakeGoodSpec() {
+  ComponentSpec component_spec;
+  component_spec.mutable_component_builder()->set_registered_name(kTestBuilder);
+  component_spec.mutable_backend()->set_registered_name(kTestBackend);
+  component_spec.mutable_transition_system()->set_registered_name("heads");
+  component_spec.mutable_network_unit()->set_registered_name(
+      "some.path.to.MstSolverNetwork");
+  LinkedFeatureChannel *link = component_spec.add_linked_feature();
+  link->set_source_component(kPreviousComponentName);
+  link->set_source_layer(kAdjacencyLayerName);
+  return component_spec;
+}
+
+class MstSolverComponentBaseTest : public NetworkTestBase {
+ protected:
+  // Initializes a head selection component from the |component_spec| and sets
+  // |heads| to the extracted head indices.  Returs non-OK on error.
+  tensorflow::Status Run(const ComponentSpec &component_spec,
+                         std::vector<int> *heads) {
+    AddComponent(kPreviousComponentName);
+    AddPairwiseLayer(kAdjacencyLayerName, 1);
+    AddPairwiseLayer(kBadDimLayerName, 2);
+
+    BasicMstSolverComponent component;
+    TF_RETURN_IF_ERROR(component.Initialize(component_spec, &variable_store_,
+                                            &network_state_manager_,
+                                            &extension_manager_));
+
+    network_states_.Reset(&network_state_manager_);
+    StartComponent(kNumSteps);
+
+    // Fill the |kRootIndex|'th column of the adjacency matrix with higher
+    // scores, so all tokens select it as head.  The |kRootIndex|'th token
+    // itself is a self-loop, so it becomes a root.
+    MutableMatrix<float> adjacency =
+        GetPairwiseLayer(kPreviousComponentName, kAdjacencyLayerName);
+    for (size_t target = 0; target < kNumSteps; ++target) {
+      for (size_t source = 0; source < kNumSteps; ++source) {
+        adjacency.row(target)[source] = source == kRootIndex ? 1.0 : 0.0;
+      }
+    }
+
+    session_state_.extensions.Reset(&extension_manager_);
+    tensorflow::gtl::ArraySlice<MstSolverComponentBase::Index> argmax;
+    TF_RETURN_IF_ERROR(component.ComputeHeads(&session_state_, &argmax));
+    heads->assign(argmax.begin(), argmax.end());
+
+    return tensorflow::Status::OK();
+  }
+};
+
+// Tests that the expected heads are produced for a good spec.
+TEST_F(MstSolverComponentBaseTest, RunsGoodSpec) {
+  std::vector<int> heads;
+  TF_ASSERT_OK(Run(MakeGoodSpec(), &heads));
+
+  const std::vector<int> expected_heads(kNumSteps, kRootIndex);
+  EXPECT_EQ(heads, expected_heads);
+}
+
+// Tests that a layer with the wrong dimension is rejected
+TEST_F(MstSolverComponentBaseTest, WrongDimension) {
+  ComponentSpec component_spec = MakeGoodSpec();
+  component_spec.mutable_linked_feature(0)->set_source_layer(kBadDimLayerName);
+
+  std::vector<int> heads;
+  EXPECT_THAT(Run(component_spec, &heads),
+              test::IsErrorWithSubstr(
+                  "Adjacency matrix has dimension 2 but expected 1"));
+}
+
+// Tests that the component is always dis-preferred.
+TEST_F(MstSolverComponentBaseTest, NotPreferred) {
+  BasicMstSolverComponent component;
+  EXPECT_FALSE(component.PreferredTo(component));
+}
+
+// Tests that the good spec is supported.
+TEST_F(MstSolverComponentBaseTest, SupportsGoodSpec) {
+  ComponentSpec component_spec = MakeGoodSpec();
+
+  BasicMstSolverComponent component;
+  EXPECT_TRUE(component.Supports(component_spec, kTestBuilder));
+}
+
+// Tests that various bad specs are rejected.
+TEST_F(MstSolverComponentBaseTest, RejectsBadSpecs) {
+  ComponentSpec component_spec = MakeGoodSpec();
+  BasicMstSolverComponent component;
+  EXPECT_FALSE(component.Supports(component_spec, "bad"));
+
+  component_spec = MakeGoodSpec();
+  component_spec.mutable_backend()->set_registered_name("bad");
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+
+  component_spec = MakeGoodSpec();
+  component_spec.mutable_transition_system()->set_registered_name("bad");
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+
+  component_spec = MakeGoodSpec();
+  component_spec.mutable_network_unit()->set_registered_name("bad");
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+
+  component_spec = MakeGoodSpec();
+  component_spec.add_fixed_feature();
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+
+  component_spec = MakeGoodSpec();
+  component_spec.add_linked_feature();
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+
+  component_spec = MakeGoodSpec();
+  component_spec.clear_linked_feature();
+  EXPECT_FALSE(component.Supports(component_spec, kTestBuilder));
+}
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/multiarch.bzl

+# Copyright 2017 Google Inc. 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.
+# =============================================================================
+
+"""Build extension rules for handling multiple target architectures."""
+
+# Build configs for specific CPU architectures. Each entry specified
+# additional copts and tags.
+# TODO(googleuser): Figure out a workaround for the shift
+# instructions, and look for any other unsupported instructions.
+MULTIARCH_CONFIGS = {
+    "generic": {
+        "copts": [],
+        "tags": [],
+    },
+    "avx": {
+        "copts": [
+            "-msse4.2",
+        ],
+        "tags": [],
+    },
+    "avx2fma": {
+        "copts": [
+            "-msse4.2",
+            "-mavx",
+            "-mavx2",
+            "-mfma",
+        ],
+        "tags": [
+            "local",
+            "manual",
+        ],
+    },
+}
+
+# List of targets which are built for multiple architectures. These
+# dependencies in dragnn_cc_* build rules are replaced with one with the
+# appropriate suffix, e.g. _multiarch_generic
+MULTIARCH_TARGETS = [
+    "//dragnn/runtime:biaffine_digraph_component",
+    "//dragnn/runtime:bulk_dynamic_component",
+    "//dragnn/runtime:bulk_feed_forward_network",
+    "//dragnn/runtime:bulk_lstm_network",
+    "//dragnn/runtime:feed_forward_network",
+    "//dragnn/runtime:feed_forward_network_kernel",
+    "//dragnn/runtime:feed_forward_network_layer",
+    "//dragnn/runtime:fixed_embeddings",
+    "//dragnn/runtime:linked_embeddings",
+    "//dragnn/runtime:lstm_network",
+    "//dragnn/runtime:lstm_network_kernel",
+    "//dragnn/runtime:network_unit_base",
+    "//dragnn/runtime:sequence_bulk_dynamic_component",
+    "//dragnn/runtime:sequence_features",
+    "//dragnn/runtime:sequence_links",
+    "//dragnn/runtime:sequence_model",
+    "//dragnn/runtime/lstm_cell:cell_function",
+    "//dragnn/runtime/lstm_cell:test_helpers",
+    "//dragnn/runtime/myelin:myelin_dynamic_component",
+    "//dragnn/runtime/myelin:myelin_dynamic_component_base",
+    "//dragnn/runtime/myelin:sequence_myelin_dynamic_component",
+    "//dragnn/runtime/xla:sequence_xla_dynamic_component_mixin",
+    "//dragnn/runtime/xla:testdata_simple_component_library",
+    "//dragnn/runtime/xla:xla_aot_dynamic_component",
+    "//dragnn/runtime/xla:xla_dynamic_component",
+    "//dragnn/runtime/xla:xla_dynamic_component_base",
+]
+
+def multiarch_name(target_name, arch_name):
+  """Generates the multiarch version of |target_name| given |arch_name|."""
+  return target_name + '_multiarch_' + arch_name
+
+def _is_multiarch(target):
+  """Returns true if |target| is designated as a multiarch target."""
+  return (target in MULTIARCH_TARGETS or
+          ('//' + native.package_name() + target) in MULTIARCH_TARGETS)
+
+def _dragnn_cc_multiarch_target(native_rule = None,
+                                name = '',
+                                target_arch = None,
+                                target_suffix = '',
+                                copts = [],
+                                deps = [],
+                                tags = [],
+                                opts_self = False,
+                                deps_transformer = None,
+                                **kwargs):
+  """Generates a target for multiple architectures.
+
+  Using the |native_rule| (e.g. cc_library) to create a set of targets for
+  all CPU architectures listed in MULTIARCH_CONFIGS, with added suffixes
+  that designate the architecture.
+
+  When |target_arch| is set, then only that single target is generated,
+  and the name of the target is unchanged (no suffix is added).
+
+  When |opts_self| is true, then the 'copts' entry in MULTIARCH_CONFIGS
+  is additionally used to build this target.
+
+  The 'tags' entry in MULTIARCH_CONFIGS are included in the build tags.
+
+  Args:
+    native_rule: The build rule used for all generated targets
+    name: The original name of the build rule (without any suffix).
+    target_arch: When set, only this architecture is targeted.
+    target_suffix: Additional suffix to add after the architecture.
+    copts: The original compilation options for this target.
+    deps: The original dependencies for this target.
+    tags: The original build tags for this target.
+    opts_self: When true, additional copts are included.
+    deps_transformer: When set, a function to apply to the multiarch deps.
+    **kwargs: Additional args passed along to the build rule.
+  """
+  # Determine set of target architectures based on |target_arch|.
+  if target_arch:
+    if target_arch in MULTIARCH_CONFIGS:
+      arch_items = [(target_arch, MULTIARCH_CONFIGS[target_arch])]
+    else:
+      fail('Unknown target_arch value: ' + target_arch)
+  else:
+    arch_items = MULTIARCH_CONFIGS.items()
+
+  # There is one target for each architecture in |arch_items|.
+  for arch, arch_config in arch_items:
+    # Transform the multi-arch deps.
+    multiarch_deps = [multiarch_name(dep, arch) if _is_multiarch(dep) else dep
+                      for dep in deps]
+    if deps_transformer:
+      multiarch_deps = deps_transformer(multiarch_deps)
+
+    native_rule(
+        name = (name if target_arch else multiarch_name(name, arch)) + target_suffix,
+        copts = copts + arch_config['copts'] if opts_self else copts,
+        deps = multiarch_deps,
+        tags = tags + arch_config['tags'],
+        **kwargs)
+
+def _dragnn_cc_multiarch_test_target(name = None,
+                                     target_arch = None,
+                                     **kwargs):
+  """Test target wrapper which puts arch name before '_test'."""
+  test_suffix = '_test'
+  has_test_suffix = name.endswith(test_suffix)
+
+  # Keeps _test at the end of the target name.
+  test_name = name[:-len(test_suffix)] if has_test_suffix else name
+  target_suffix = test_suffix if has_test_suffix else ''
+
+  _dragnn_cc_multiarch_target(native_rule = native.cc_test,
+                              name = test_name,
+                              target_arch = target_arch,
+                              target_suffix = target_suffix,
+                              **kwargs)
+
+  # When |target_arch| is set, the resulting test is named |name|.  Otherwise,
+  # tests with arch-specific names are generated, and for convenience we add a
+  # test_suite named |name| that runs the generic version of the test.
+  if not target_arch:
+    native.test_suite(
+        name = name,
+        tests = [multiarch_name(test_name, 'generic') + target_suffix])
+
+def dragnn_cc_multiarch_library(**kwargs):
+  """Similar to cc_library, but creates multiple architecture targets."""
+
+  _dragnn_cc_multiarch_target(native_rule = native.cc_library,
+                              **kwargs)
+
+def dragnn_cc_multiarch_test(**kwargs):
+  """Similar to cc_test, but creates multiple architecture targets."""
+  _dragnn_cc_multiarch_test_target(**kwargs)
+
+def dragnn_cc_multiarch_binary(**kwargs):
+  """Similar to cc_binary, but creates multiple architecture targets."""
+  _dragnn_cc_multiarch_target(native_rule = native.cc_binary,
+                              **kwargs)
+
+def dragnn_cc_library(target_arch = 'generic', **kwargs):
+  """Similar to cc_library, but targets one specific architecture."""
+  _dragnn_cc_multiarch_target(native_rule = native.cc_library,
+                              target_arch = target_arch,
+                              **kwargs)
+
+def dragnn_cc_test(target_arch = 'generic', **kwargs):
+  """Similar to cc_test, but targets one specific architecture."""
+  _dragnn_cc_multiarch_test_target(target_arch = target_arch,
+                                   **kwargs)
+
+def dragnn_cc_binary(target_arch = 'generic', **kwargs):
+  """Similar to cc_binary, but targets one specific architecture."""
+  _dragnn_cc_multiarch_target(native_rule = native.cc_binary,
+                              target_arch = target_arch,
+                              **kwargs)

research/syntaxnet/dragnn/runtime/myelin/BUILD

+package(default_visibility = ["//visibility:public"])
+
+load(
+    ":build_defs.bzl",
+    "dragnn_myelin_cc_library",
+    "dragnn_myelin_cc_test",
+    "dragnn_myelin_cc_multiarch_library",
+    "dragnn_myelin_cc_multiarch_test",
+)
+
+test_suite(name = "all_tests")
+
+filegroup(
+    name = "test_myelination_output",
+    srcs = glob(["testdata/myelination_output/**"]),
+)
+
+cc_library(
+    name = "attr_value_utils",
+    srcs = ["attr_value_utils.cc"],
+    hdrs = ["attr_value_utils.h"],
+    deps = [
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:framework_headers_lib",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:protos_all_cc",
+    ],
+)
+
+cc_test(
+    name = "attr_value_utils_test",
+    size = "small",
+    srcs = ["attr_value_utils_test.cc"],
+    deps = [
+        ":attr_value_utils",
+        "//dragnn/core/test:generic",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+    ],
+)
+
+dragnn_myelin_cc_library(
+    name = "myelin_cell_converter",
+    srcs = ["myelin_cell_converter.cc"],
+    hdrs = ["myelin_cell_converter.h"],
+    deps = [
+        ":attr_value_utils",
+        "//dragnn/protos:export_proto_cc",
+        "//dragnn/runtime:trained_model",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:framework_headers_lib",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:protos_all_cc",
+    ],
+)
+
+dragnn_myelin_cc_test(
+    name = "myelin_cell_converter_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["myelin_cell_converter_test.cc"],
+    data = ["//dragnn/runtime:test_rnn_tagger"],
+    deps = [
+        ":myelin_cell_converter",
+        ":myelin_spec_utils",
+        "//dragnn/components/syntaxnet:syntaxnet_component",
+        "//dragnn/core/test:generic",
+        "//dragnn/runtime:alignment",
+        "//dragnn/runtime:trained_model",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@sling//sling/myelin:compute",
+        "@sling//sling/myelin:flow",
+        "@sling//sling/myelin:graph",
+    ],
+)
+
+dragnn_myelin_cc_library(
+    name = "myelin_library",
+    srcs = ["myelin_library.cc"],
+    hdrs = ["myelin_library.h"],
+    deps = [
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/myelin:flow",
+    ],
+)
+
+dragnn_myelin_cc_test(
+    name = "myelin_library_test",
+    size = "small",
+    srcs = ["myelin_library_test.cc"],
+    deps = [
+        ":myelin_library",
+        "//syntaxnet:test_main",
+        "@org_tensorflow//tensorflow/core:test",
+    ],
+)
+
+dragnn_myelin_cc_library(
+    name = "myelin_spec_utils",
+    srcs = ["myelin_spec_utils.cc"],
+    hdrs = ["myelin_spec_utils.h"],
+    deps = [
+        ":myelin_library",
+        "//dragnn/protos:spec_proto_cc",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/base",
+        "@sling//sling/file",
+        "@sling//sling/myelin:compute",
+        "@sling//sling/myelin:flow",
+        "@sling//sling/myelin/kernel:tensorflow",
+    ],
+)
+
+dragnn_myelin_cc_test(
+    name = "myelin_spec_utils_test",
+    size = "small",
+    srcs = ["myelin_spec_utils_test.cc"],
+    deps = [
+        ":myelin_spec_utils",
+        "//dragnn/core/test:generic",
+        "//dragnn/protos:spec_proto_cc",
+        "//syntaxnet:base",
+        "//syntaxnet:test_main",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@sling//sling/file",
+        "@sling//sling/file:posix",
+        "@sling//sling/myelin:compute",
+        "@sling//sling/myelin:flow",
+    ],
+)
+
+dragnn_myelin_cc_library(
+    name = "myelin_tracing",
+    srcs = ["myelin_tracing.cc"],
+    hdrs = ["myelin_tracing.h"],
+    deps = [
+        "//dragnn/protos:cell_trace_proto_cc",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/myelin:compute",
+    ],
+)
+
+dragnn_myelin_cc_test(
+    name = "myelin_tracing_test",
+    size = "small",
+    srcs = ["myelin_tracing_test.cc"],
+    deps = [
+        ":myelin_spec_utils",
+        ":myelin_tracing",
+        "//dragnn/core/test:generic",
+        "//dragnn/protos:cell_trace_proto_cc",
+        "//dragnn/runtime/test:helpers",
+        "//syntaxnet:base",
+        "//syntaxnet:test_main",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@sling//sling/myelin:compute",
+        "@sling//sling/myelin:flow",
+    ],
+)
+
+dragnn_myelin_cc_multiarch_library(
+    name = "myelin_dynamic_component_base",
+    srcs = ["myelin_dynamic_component_base.cc"],
+    hdrs = ["myelin_dynamic_component_base.h"],
+    deps = [
+        ":myelin_spec_utils",
+        ":myelin_tracing",
+        "//dragnn/protos:cell_trace_proto_cc",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/protos:trace_proto_cc",
+        "//dragnn/runtime:alignment",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:extensions",
+        "//dragnn/runtime:fixed_embeddings",
+        "//dragnn/runtime:linked_embeddings",
+        "//dragnn/runtime:network_states",
+        "//dragnn/runtime:session_state",
+        "//dragnn/runtime:transition_system_traits",
+        "//dragnn/runtime:variable_store",
+        "//dragnn/runtime/math:types",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/myelin:compute",
+        "@sling//sling/myelin:flow",
+    ],
+)
+
+dragnn_myelin_cc_multiarch_library(
+    name = "myelin_dynamic_component",
+    srcs = ["myelin_dynamic_component.cc"],
+    deps = [
+        ":myelin_dynamic_component_base",
+        "//dragnn/core:compute_session",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/protos:trace_proto_cc",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:fixed_embeddings",
+        "//dragnn/runtime:linked_embeddings",
+        "//dragnn/runtime:network_states",
+        "//dragnn/runtime:session_state",
+        "//dragnn/runtime/math:types",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/myelin:compute",
+    ],
+    alwayslink = 1,
+)
+
+dragnn_myelin_cc_multiarch_test(
+    name = "myelin_dynamic_component_test",
+    size = "small",
+    srcs = ["myelin_dynamic_component_test.cc"],
+    deps = [
+        ":myelin_dynamic_component",
+        ":myelin_spec_utils",
+        "//dragnn/core/test:generic",
+        "//dragnn/protos:cell_trace_proto_cc",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/protos:trace_proto_cc",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:extensions",
+        "//dragnn/runtime/test:network_test_base",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@sling//sling/file",
+        "@sling//sling/file:posix",
+        "@sling//sling/myelin:flow",
+    ],
+)
+
+dragnn_myelin_cc_library(
+    name = "myelination",
+    srcs = ["myelination.cc"],
+    hdrs = ["myelination.h"],
+    deps = [
+        ":myelin_cell_converter",
+        ":myelin_spec_utils",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:trained_model",
+        "//syntaxnet:base",
+        "//syntaxnet:registry",
+        "@org_tensorflow//tensorflow/core:lib",
+    ],
+)
+
+dragnn_myelin_cc_test(
+    name = "myelination_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["myelination_test.cc"],
+    data = [
+        ":test_myelination_output",
+        "//dragnn/runtime:test_rnn_tagger",
+    ],
+    deps = [
+        ":myelin_spec_utils",
+        ":myelination",
+        "//dragnn/components/syntaxnet:syntaxnet_component",
+        "//dragnn/core/test:generic",
+        "//dragnn/protos:spec_proto_cc",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+    ],
+)
+
+dragnn_myelin_cc_multiarch_library(
+    name = "sequence_myelin_dynamic_component",
+    srcs = ["sequence_myelin_dynamic_component.cc"],
+    deps = [
+        ":myelin_dynamic_component_base",
+        "//dragnn/core:compute_session",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/protos:trace_proto_cc",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:extensions",
+        "//dragnn/runtime:network_states",
+        "//dragnn/runtime:sequence_features",
+        "//dragnn/runtime:sequence_links",
+        "//dragnn/runtime:sequence_model",
+        "//dragnn/runtime:session_state",
+        "//dragnn/runtime:variable_store",
+        "//dragnn/runtime/math:types",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@sling//sling/myelin:compute",
+    ],
+    alwayslink = 1,
+)
+
+dragnn_myelin_cc_multiarch_test(
+    name = "sequence_myelin_dynamic_component_test",
+    size = "small",
+    srcs = ["sequence_myelin_dynamic_component_test.cc"],
+    deps = [
+        ":myelin_spec_utils",
+        ":sequence_myelin_dynamic_component",
+        "//dragnn/core:compute_session",
+        "//dragnn/core:input_batch_cache",
+        "//dragnn/core/test:generic",
+        "//dragnn/protos:spec_proto_cc",
+        "//dragnn/protos:trace_proto_cc",
+        "//dragnn/runtime:component",
+        "//dragnn/runtime:extensions",
+        "//dragnn/runtime:network_states",
+        "//dragnn/runtime:sequence_backend",
+        "//dragnn/runtime:sequence_extractor",
+        "//dragnn/runtime:sequence_linker",
+        "//dragnn/runtime:sequence_predictor",
+        "//dragnn/runtime:session_state",
+        "//dragnn/runtime:variable_store",
+        "//dragnn/runtime/math:types",
+        "//dragnn/runtime/test:network_test_base",
+        "//syntaxnet:base",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@sling//sling/file",
+        "@sling//sling/file:posix",
+        "@sling//sling/myelin:flow",
+    ],
+)

research/syntaxnet/dragnn/runtime/myelin/attr_value_utils.cc

+// Copyright 2017 Google Inc. 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.
+// =============================================================================
+
+// Implementation note: This file contains branched versions of functions from
+// tensorflow/core/framework/attr_value_util.cc.  These functions are branched
+// to prevent changes in their behavior from impacting the Myelin conversion.
+
+#include "dragnn/runtime/myelin/attr_value_utils.h"
+
+#include <algorithm>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+using ::tensorflow::AttrValue;
+using ::tensorflow::NameAttrList;
+using ::tensorflow::PartialTensorShape;
+using ::tensorflow::StringPiece;
+using ::tensorflow::Tensor;
+using ::tensorflow::TensorProto;
+using ::tensorflow::TensorShape;
+namespace strings = ::tensorflow::strings;
+namespace str_util = ::tensorflow::str_util;
+
+string SummarizeString(const string &str) {
+  string escaped = str_util::CEscape(str);
+
+  // If the string is long, replace the middle with ellipses.
+  constexpr int kMaxStringSummarySize = 80;
+  if (escaped.size() >= kMaxStringSummarySize) {
+    StringPiece prefix(escaped);
+    StringPiece suffix = prefix;
+    prefix.remove_suffix(escaped.size() - 10);
+    suffix.remove_prefix(escaped.size() - 10);
+    return strings::StrCat("\"", prefix, "...", suffix, "\"");
+  } else {
+    return strings::StrCat("\"", escaped, "\"");
+  }
+}
+
+string SummarizeTensor(const TensorProto &tensor_proto) {
+  Tensor t;
+  if (!t.FromProto(tensor_proto)) return "<Invalid TensorProto>";
+  return t.DebugString();
+}
+
+string SummarizeFunc(const NameAttrList &func) {
+  std::vector<string> entries;
+  for (auto p : func.attr()) {
+    entries.push_back(
+        strings::StrCat(p.first, "=", AttrValueToString(p.second)));
+  }
+  std::sort(entries.begin(), entries.end());
+  return strings::StrCat(func.name(), "[", str_util::Join(entries, ", "), "]");
+}
+
+}  // namespace
+
+string AttrValueToString(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 DataType_Name(attr_value.type());
+    case AttrValue::kShape:
+      return PartialTensorShape::DebugString(attr_value.shape());
+    case AttrValue::kTensor:
+      return SummarizeTensor(attr_value.tensor());
+    case AttrValue::kList: {
+      std::vector<string> pieces;
+      if (attr_value.list().s_size() > 0) {
+        for (int i = 0; i < attr_value.list().s_size(); ++i) {
+          pieces.push_back(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) {
+          pieces.push_back(strings::StrCat(attr_value.list().i(i)));
+        }
+      } else if (attr_value.list().f_size() > 0) {
+        for (int i = 0; i < attr_value.list().f_size(); ++i) {
+          pieces.push_back(strings::StrCat(attr_value.list().f(i)));
+        }
+      } else if (attr_value.list().b_size() > 0) {
+        for (int i = 0; i < attr_value.list().b_size(); ++i) {
+          pieces.push_back(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) {
+          pieces.push_back(DataType_Name(attr_value.list().type(i)));
+        }
+      } else if (attr_value.list().shape_size() > 0) {
+        for (int i = 0; i < attr_value.list().shape_size(); ++i) {
+          pieces.push_back(
+              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) {
+          pieces.push_back(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) {
+          pieces.push_back(SummarizeFunc(attr_value.list().func(i)));
+        }
+      }
+      return strings::StrCat("[", str_util::Join(pieces, ", "), "]");
+    }
+    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
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/attr_value_utils.h

+// Copyright 2017 Google Inc. 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.
+// =============================================================================
+
+// Utils for working with tensorflow.AttrValue protos.
+
+#ifndef DRAGNN_RUNTIME_MYELIN_ATTR_VALUE_UTILS_H_
+#define DRAGNN_RUNTIME_MYELIN_ATTR_VALUE_UTILS_H_
+
+#include <string>
+
+#include "syntaxnet/base.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// Returns a string representation of the |attr_value|.  This is similar to
+// tensorflow::SummarizeAttrValue(), but never elides or abbreviates.
+string AttrValueToString(const tensorflow::AttrValue &attr_value);
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MYELIN_ATTR_VALUE_UTILS_H_

research/syntaxnet/dragnn/runtime/myelin/attr_value_utils_test.cc

+// Copyright 2017 Google Inc. 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.
+// =============================================================================
+
+// NB: These tests don't assert on dtypes, shapes, or tensors, because those are
+// just calls to TF library functions.  (I.e., don't test someone else's API).
+
+#include "dragnn/runtime/myelin/attr_value_utils.h"
+
+#include <string>
+
+#include "dragnn/core/test/generic.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+// Tests that singular attributes are stringified correctly.
+TEST(AttrValueToStringTest, Singular) {
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.set_s("foo");
+    EXPECT_EQ(AttrValueToString(attr_value), "\"foo\"");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.set_i(123);
+    EXPECT_EQ(AttrValueToString(attr_value), "123");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.set_f(-1.5);
+    EXPECT_EQ(AttrValueToString(attr_value), "-1.5");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.set_b(false);
+    EXPECT_EQ(AttrValueToString(attr_value), "false");
+    attr_value.set_b(true);
+    EXPECT_EQ(AttrValueToString(attr_value), "true");
+  }
+}
+
+// Tests that list attributes are stringified correctly.
+TEST(AttrValueToStringTest, List) {
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.mutable_list()->add_s("foo");
+    attr_value.mutable_list()->add_s("bar");
+    attr_value.mutable_list()->add_s("baz");
+    EXPECT_EQ(AttrValueToString(attr_value), "[\"foo\", \"bar\", \"baz\"]");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.mutable_list()->add_i(123);
+    attr_value.mutable_list()->add_i(-45);
+    attr_value.mutable_list()->add_i(6789);
+    EXPECT_EQ(AttrValueToString(attr_value), "[123, -45, 6789]");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.mutable_list()->add_f(-1.5);
+    attr_value.mutable_list()->add_f(0.25);
+    attr_value.mutable_list()->add_f(3.5);
+    EXPECT_EQ(AttrValueToString(attr_value), "[-1.5, 0.25, 3.5]");
+  }
+
+  {
+    tensorflow::AttrValue attr_value;
+    attr_value.mutable_list()->add_b(false);
+    attr_value.mutable_list()->add_b(true);
+    attr_value.mutable_list()->add_b(false);
+    EXPECT_EQ(AttrValueToString(attr_value), "[false, true, false]");
+  }
+}
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/build_defs.bzl

+# Copyright 2017 Google Inc. 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.
+# =============================================================================
+
+"""Build rules that restrict Myelin to supported environments.
+
+As of this writing, SLING requires Linux and x86-64:
+https://github.com/google/sling/blob/master/README.md#building
+
+The technique used here is to replace the hdrs, srcs, and deps with appropriate
+empty content when building in an unsupported environment.
+"""
+
+load(
+    "//dragnn/runtime:multiarch.bzl",
+    "dragnn_cc_multiarch_library",
+    "dragnn_cc_multiarch_test",
+)
+
+def _if_supported(consequent, alternative=[]):
+  """Returns the |consequent| iff the build environment supports Myelin."""
+  return select({
+      "@org_tensorflow//tensorflow:linux_x86_64": consequent,
+      "//conditions:default": alternative,
+  })
+
+def _if_supported_test_deps(deps):
+  """Like _if_supported, but returns appropriate fallback deps for a test."""
+  return _if_supported(deps, ["//syntaxnet:test_main"])
+
+def dragnn_myelin_cc_library(hdrs=[], srcs=[], deps=[], **kwargs):
+  """Like cc_library, but reduces to a NOP in unsupported environments."""
+  native.cc_library(
+      hdrs = _if_supported(hdrs),
+      srcs = _if_supported(srcs),
+      deps = _if_supported(deps),
+      **kwargs)
+
+def dragnn_myelin_cc_test(srcs=[], deps=[], **kwargs):
+  """Like cc_test, but reduces to a NOP in unsupported environments."""
+  native.cc_test(
+      srcs = _if_supported(srcs),
+      deps = _if_supported_test_deps(deps),
+      **kwargs)
+
+# Implementation note: Bazel select()s are not resolved at the time that build
+# rules are evaluated.  If we pass _if_supported(deps) into the multi-arch build
+# rules (like we do for the native rules above), then the multi-arch rules break
+# when they attempt to iterate over the deps---at that point, the deps are an
+# unresolved select() that can't be iterated.  To get around this, we delay the
+# select() by passing _if_supported into the multi-arch rule, which will apply
+# it just before passing the deps to the native rule.
+
+def dragnn_myelin_cc_multiarch_library(hdrs=[], srcs=[], **kwargs):
+  """Multi-arch version of dragnn_myelin_cc_library."""
+  dragnn_cc_multiarch_library(
+      hdrs = _if_supported(hdrs),
+      srcs = _if_supported(srcs),
+      deps_transformer = _if_supported,
+      **kwargs)
+
+def dragnn_myelin_cc_multiarch_test(srcs=[], **kwargs):
+  """Multi-arch version of dragnn_myelin_cc_test."""
+  dragnn_cc_multiarch_test(
+      srcs = _if_supported(srcs, []),
+      deps_transformer = _if_supported_test_deps,
+      **kwargs)

research/syntaxnet/dragnn/runtime/myelin/myelin_cell_converter.cc

+// Copyright 2017 Google Inc. 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 "dragnn/runtime/myelin/myelin_cell_converter.h"
+
+#include <stddef.h>
+#include <algorithm>
+#include <limits>
+
+#include "dragnn/runtime/myelin/attr_value_utils.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/cpu_info.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+// Returns true if the |tensor_name| denotes a control dependency.
+bool IsControlDependency(const string &tensor_name) {
+  return tensor_name[0] == '^';
+}
+
+// Returns true if the |node| is a TF variable.
+bool IsVariableNode(const tensorflow::NodeDef &node) {
+  return node.op() == "VariableV2";
+}
+
+// Returns true if the |node| is a tf.constant().
+bool IsConstantNode(const tensorflow::NodeDef &node) {
+  return node.op() == "Const";
+}
+
+// Returns true if the |node| is a tf.placeholder().
+bool IsPlaceholderNode(const tensorflow::NodeDef &node) {
+  return node.op() == "Placeholder";
+}
+
+// Sets |max_value| to |value| if it is lesser.
+void UpdateMax(uint32 value, uint32 *max_value) {
+  *max_value = std::max(*max_value, value);
+}
+
+// Loads the |tensor| from the constant |node|.  On error, returns non-OK.
+tensorflow::Status GetConstantTensor(const tensorflow::NodeDef &node,
+                                     tensorflow::Tensor *tensor) {
+  DCHECK(IsConstantNode(node));
+  return tensorflow::GetNodeAttr(node, "value", tensor);
+}
+
+// Loads the |shape| from the placeholder |node|.  On error, returns non-OK.
+tensorflow::Status GetPlaceholderShape(const tensorflow::NodeDef &node,
+                                       tensorflow::TensorShape *shape) {
+  DCHECK(IsPlaceholderNode(node));
+  return tensorflow::GetNodeAttr(node, "shape", shape);
+}
+
+// Returns the dtype string associated with the |node|, or an empty string if it
+// cannot be inferred.
+string GetDType(const tensorflow::NodeDef &node) {
+  tensorflow::DataType dtype;
+  tensorflow::Status status = tensorflow::GetNodeAttr(node, "T", &dtype);
+  if (!status.ok()) status = tensorflow::GetNodeAttr(node, "dtype", &dtype);
+  if (status.ok()) return tensorflow::DataTypeString(dtype);
+  return string();
+}
+
+// Modifies the |dtype| into a reference type.
+void MarkAsReferenceDType(string *dtype) {
+  DCHECK_NE((*dtype)[0], '&');
+  *dtype = tensorflow::strings::StrCat("&", *dtype);
+}
+
+// Loads the CellSubgraphSpec for the component named |component_name| from the
+// |trained_model| into the |spec|.  On error, returns non-OK.
+tensorflow::Status LoadCellSubgraphSpec(const string &component_name,
+                                        const TrainedModel &trained_model,
+                                        CellSubgraphSpec *spec) {
+  const string tensor_name =
+      tensorflow::strings::StrCat(component_name, "/EXPORT/CellSubgraphSpec");
+  tensorflow::Tensor tensor;
+  TF_RETURN_IF_ERROR(trained_model.EvaluateTensor(tensor_name, &tensor));
+
+  if (!spec->ParseFromString(tensor.scalar<string>()())) {
+    return tensorflow::errors::InvalidArgument(
+        "Failed to parse CellSubgraphSpec for component ", component_name);
+  }
+
+  VLOG(1) << tensor_name << " = \n" << spec->DebugString();
+  return tensorflow::Status::OK();
+}
+
+}  // namespace
+
+// Writer for incrementally building a Flow file.
+// https://github.com/google/sling/tree/master/myelin#flow-file-format
+
+class MyelinCellConverter::Writer {
+ public:
+  // TODO(googleuser): Add templated Write() methods and coerce typed data into
+  // little-endian format, so this doesn't need to run on a little-endian CPU.
+  static_assert(tensorflow::port::kLittleEndian,
+                "Flow files must be written in little-endian format");
+
+  // Creates a writer that overwrites |flow|.
+  explicit Writer(string *flow) : flow_(CHECK_NOTNULL(flow)) {
+    flow_->clear();
+    Write("flow", 4);  // magic number
+    WriteInt32(4);  // version
+  }
+
+  // Appends [|data|,|data|+|size|) to the Flow file.
+  void Write(const void *data, size_t size) {
+    flow_->append(reinterpret_cast<const char *>(data), size);
+  }
+
+  // Appends the |value| to the Flow file.
+  void WriteInt32(int32 value) { Write(&value, sizeof(int32)); }
+  void WriteUint64(uint64 value) { Write(&value, sizeof(uint64)); }
+
+  // Writes the |str| to the Flow file as a length-prefixed string.
+  void WriteString(const string &str) {
+    DCHECK_LE(str.size(), std::numeric_limits<int32>::max());
+    WriteInt32(str.size());
+    Write(str.data(), str.size());
+  }
+
+ private:
+  // Flow file content.
+  string *const flow_;
+};
+
+tensorflow::Status MyelinCellConverter::Convert(
+    const string &component_name, const TrainedModel &trained_model,
+    string *flow) {
+  return MyelinCellConverter().ConvertImpl(component_name, trained_model, flow);
+}
+
+tensorflow::Status MyelinCellConverter::ConvertImpl(
+    const string &component_name, const TrainedModel &trained_model,
+    string *flow) {
+  component_name_ = component_name;
+  trained_model_ = &trained_model;
+
+  CellSubgraphSpec spec;
+  TF_RETURN_IF_ERROR(
+      LoadCellSubgraphSpec(component_name_, *trained_model_, &spec));
+  TF_RETURN_IF_ERROR(BuildInputsAndOutputs(spec));
+  TF_RETURN_IF_ERROR(BuildOperations());
+
+  Writer writer(flow);
+  TF_RETURN_IF_ERROR(WriteVariables(&writer));
+  WriteOperations(&writer);
+  WriteFunctions(&writer);
+  WriteConnectors(&writer);
+  WriteBlobs(&writer);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinCellConverter::BuildInputsAndOutputs(
+    const CellSubgraphSpec &spec) {
+  std::set<string> unique_input_names;
+  for (const CellSubgraphSpec::Input &input : spec.input()) {
+    if (!unique_input_names.insert(input.name()).second) {
+      return tensorflow::errors::InvalidArgument(
+          "Duplicate input name { ", input.ShortDebugString(), " }");
+    }
+
+    TensorId tensor_id;
+    TF_RETURN_IF_ERROR(ParseTensorId(input.tensor(), &tensor_id));
+
+    if (inputs_.find(tensor_id) != inputs_.end()) {
+      return tensorflow::errors::InvalidArgument(
+          "Duplicate input variable { ", input.ShortDebugString(),
+          " }; currently has name '", inputs_[tensor_id], "'");
+    }
+
+    inputs_[tensor_id] = input.name();
+  }
+
+  std::set<string> unique_output_names;
+  for (const CellSubgraphSpec::Output &output : spec.output()) {
+    if (!unique_output_names.insert(output.name()).second) {
+      return tensorflow::errors::InvalidArgument(
+          "Duplicate output name { ", output.ShortDebugString(), " }");
+    }
+
+    TensorId tensor_id;
+    TF_RETURN_IF_ERROR(ParseTensorId(output.tensor(), &tensor_id));
+
+    outputs_[tensor_id].insert(output.name());
+  }
+
+  // Check that recurrent inputs match the name of an output.
+  for (const CellSubgraphSpec::Input &input : spec.input()) {
+    if (input.type() != CellSubgraphSpec::Input::TYPE_RECURRENT) continue;
+
+    if (unique_output_names.find(input.name()) == unique_output_names.end()) {
+      return tensorflow::errors::InvalidArgument(
+          "Recurrent input does not match any output { ",
+          input.ShortDebugString(), " }");
+    }
+  }
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinCellConverter::BuildOperations() {
+  // Extract sets of input and output node names.
+  std::set<string> input_node_names;
+  std::set<string> output_node_names;
+  for (const auto &it : inputs_) input_node_names.insert(it.first.first);
+  for (const auto &it : outputs_) output_node_names.insert(it.first.first);
+
+  // Set of nodes that have already been visited by the DFS.
+  std::set<string> visited;
+
+  // DFS backwards from output nodes to input nodes and collect operations.
+  std::vector<string> stack(output_node_names.begin(), output_node_names.end());
+  while (!stack.empty()) {
+    const string name = stack.back();
+    stack.pop_back();
+    if (!visited.insert(name).second) continue;  // already visited; skip
+
+    const tensorflow::NodeDef *node = nullptr;
+    TF_RETURN_IF_ERROR(trained_model_->LookupNode(name, &node));
+
+    Operation &operation = operations_[name];
+    if (operation.node != nullptr && operation.node != node) {
+      return tensorflow::errors::Internal("Inconsistent nodes for operation ",
+                                          name, " (", operation.node->name(),
+                                          " vs ", node->name());
+    }
+    operation.node = node;
+
+    // Function inputs bound the search; don't expand them.
+    if (input_node_names.find(name) != input_node_names.end()) continue;
+
+    // Expand (non-control) inputs.
+    for (const string &input_name : node->input()) {
+      if (IsControlDependency(input_name)) continue;
+      VLOG(1) << name << " has input " << input_name;
+
+      TensorId tensor_id;
+      TF_RETURN_IF_ERROR(ParseTensorId(input_name, &tensor_id));
+      stack.push_back(tensor_id.first);
+
+      // Add the input tensor and register the output index on the input op.
+      operation.inputs.push_back(AsVariableName(tensor_id));
+      UpdateMax(tensor_id.second + 1,
+                &operations_[tensor_id.first].num_outputs);
+    }
+  }
+
+  // Register output indices for the |outputs_|; the DFS does not cover these.
+  for (const auto &it : outputs_) {
+    const TensorId &tensor_id = it.first;
+    UpdateMax(tensor_id.second + 1, &operations_[tensor_id.first].num_outputs);
+  }
+
+  // Sanity check: All operations must have nodes and outputs.
+  for (const auto &it : operations_) {
+    const Operation &operation = it.second;
+    DCHECK(operation.node != nullptr);
+    DCHECK_GT(operation.num_outputs, 0);
+  }
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinCellConverter::WriteVariables(Writer *writer) const {
+  int num_variables = 0;
+  for (const auto &it : operations_) num_variables += it.second.num_outputs;
+  writer->WriteInt32(num_variables);
+
+  for (const auto &it : operations_) {
+    const Operation &operation = it.second;
+    for (uint32 output = 0; output < operation.num_outputs; ++output) {
+      TF_RETURN_IF_ERROR(WriteVariable(*operation.node, output, writer));
+    }
+  }
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinCellConverter::WriteVariable(
+    const tensorflow::NodeDef &node, uint32 output_index,
+    Writer *writer) const {
+  const TensorId tensor_id(node.name(), output_index);
+  const string name = AsVariableName(tensor_id);
+  const std::set<string> aliases = GetAliases(tensor_id);
+
+  // Only cell inputs and outputs have aliases.
+  const bool is_cell_input_or_output = !aliases.empty();
+
+  // Treat cell inputs and outputs as references, so they can be pointed at
+  // pieces of memory managed by the DRAGNN runtime.
+  string dtype = GetDType(node);
+  if (is_cell_input_or_output) MarkAsReferenceDType(&dtype);
+
+  // Extract variable data and shape, if available.  Myelin treats a 0-element
+  // shape (e.g., [0], [1, 0, 2]) as undefined and will infer shapes for such
+  // variables, so we ensure that the shape is undefined unless explicitly set.
+  tensorflow::Tensor tensor;
+  tensorflow::TensorShape shape({0});  // undefined by default
+  if (IsConstantNode(node)) {
+    TF_RETURN_IF_ERROR(GetConstantTensor(node, &tensor));
+    shape = tensor.shape();
+  } else if (IsVariableNode(node)) {
+    TF_RETURN_IF_ERROR(trained_model_->EvaluateTensor(name, &tensor));
+    shape = tensor.shape();
+  } else if (IsPlaceholderNode(node)) {
+    TF_RETURN_IF_ERROR(GetPlaceholderShape(node, &shape));
+  }
+  const tensorflow::StringPiece data = tensor.tensor_data();
+
+  writer->WriteString(name);
+  writer->WriteInt32(aliases.size());
+  for (const string &alias : aliases) writer->WriteString(alias);
+  writer->WriteString(dtype);
+
+  writer->WriteInt32(shape.dims());
+  for (int i = 0; i < shape.dims(); ++i) writer->WriteInt32(shape.dim_size(i));
+
+  writer->WriteUint64(data.size());
+  writer->Write(data.data(), data.size());
+
+  return tensorflow::Status::OK();
+}
+
+std::set<string> MyelinCellConverter::GetAliases(
+    const TensorId &tensor_id) const {
+  std::set<string> aliases;
+
+  const auto input_it = inputs_.find(tensor_id);
+  if (input_it != inputs_.end()) {
+    const string &name = input_it->second;
+    aliases.insert(tensorflow::strings::StrCat("INPUT/", name));
+  }
+
+  const auto output_it = outputs_.find(tensor_id);
+  if (output_it != outputs_.end()) {
+    for (const string &name : output_it->second) {
+      aliases.insert(tensorflow::strings::StrCat("OUTPUT/", name));
+    }
+  }
+
+  return aliases;
+}
+
+void MyelinCellConverter::WriteOperations(Writer *writer) const {
+  writer->WriteInt32(operations_.size());
+  for (const auto &it : operations_) {
+    const Operation &operation = it.second;
+    WriteOperation(operation, writer);
+  }
+}
+
+void MyelinCellConverter::WriteOperation(const Operation &operation,
+                                         Writer *writer) const {
+  const string &name = operation.node->name();
+  const string &type = operation.node->op();
+
+  // Create one output per possible output index, in order.
+  std::vector<string> outputs;
+  for (uint32 output = 0; output < operation.num_outputs; ++output) {
+    outputs.push_back(AsVariableName(TensorId(name, output)));
+  }
+
+  // Copy the attrs to a sorted map for deterministic ordering.
+  std::map<string, tensorflow::AttrValue> attrs(operation.node->attr().begin(),
+                                                operation.node->attr().end());
+
+  writer->WriteString(name);
+  writer->WriteString(type);
+
+  writer->WriteInt32(operation.inputs.size());
+  for (const string &input : operation.inputs) writer->WriteString(input);
+
+  writer->WriteInt32(outputs.size());
+  for (const string &output : outputs) writer->WriteString(output);
+
+  writer->WriteInt32(attrs.size());
+  for (const auto &it : attrs) {
+    writer->WriteString(it.first);
+    writer->WriteString(AttrValueToString(it.second));
+  }
+}
+
+void MyelinCellConverter::WriteFunctions(Writer *writer) const {
+  writer->WriteInt32(1);
+  writer->WriteString(component_name_);
+  writer->WriteInt32(operations_.size());
+  for (const auto &it : operations_) writer->WriteString(it.first);
+}
+
+void MyelinCellConverter::WriteConnectors(Writer *writer) const {
+  writer->WriteInt32(0);
+}
+
+void MyelinCellConverter::WriteBlobs(Writer *writer) const {
+  writer->WriteInt32(0);
+}
+
+tensorflow::Status MyelinCellConverter::ParseTensorId(const string &tensor_name,
+                                                      TensorId *tensor_id) {
+  if (IsControlDependency(tensor_name)) {
+    return tensorflow::errors::InvalidArgument(
+        "Cannot parse tensor ID from control dependency '", tensor_name, "'");
+  }
+
+  const auto colon_index = tensor_name.rfind(':');
+
+  // NB: If |colon_index| is string::npos, takes the whole string as desired.
+  tensor_id->first = tensor_name.substr(0, colon_index);
+
+  if (colon_index == string::npos) {  // no colon; assume 0
+    tensor_id->second = 0;
+  } else {
+    const string output_str = tensor_name.substr(colon_index + 1);
+    if (!tensorflow::strings::safe_strtou32(output_str, &tensor_id->second)) {
+      return tensorflow::errors::InvalidArgument("Malformed tensor name ",
+                                                 tensor_name);
+    }
+  }
+
+  return tensorflow::Status::OK();
+}
+
+string MyelinCellConverter::AsVariableName(const TensorId &tensor_id) {
+  if (tensor_id.second == 0) return tensor_id.first;
+  return tensorflow::strings::StrCat(tensor_id.first, ":", tensor_id.second);
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/myelin_cell_converter.h

+// Copyright 2017 Google Inc. 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 DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_
+#define DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_
+
+#include <map>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "dragnn/protos/export.pb.h"
+#include "dragnn/runtime/trained_model.h"
+#include "syntaxnet/base.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// Converter that extracts the cell computation from a DRAGNN component and
+// writes it as a Myelin Flow.
+
+//
+// The trained model that contains the DRAGNN component must also contain a
+// CellSubgraphSpec proto embedded into the TF graph as a specifically-named
+// constant node (see runtime_support.py).  The CellSubgraphSpec defines the
+// boundaries of the cell comptation.
+//
+// The converted Myelin Flow contains a single function that runs the cell and
+// is named after the component.  The function inputs and outputs are reference
+// variables, so they can be pointed at externally-managed pieces of memory,
+// provided sufficient size and alignment.  The function inputs and outputs are
+// marked with special aliases, namely:
+//   INPUT/<CellSubgraphSpec.Input.name>
+//   OUTPUT/<CellSubgraphSpec.Output.name>
+class MyelinCellConverter {
+ public:
+  // Extracts the cell of the DRAGNN component named |component_name| from the
+  // |trained_model| and overwrites the |flow| with an equivalent Myelin Flow.
+  // The |flow| file output is deterministic given identical inputs.  On error,
+  // returns non-OK.
+  static tensorflow::Status Convert(const string &component_name,
+                                    const TrainedModel &trained_model,
+                                    string *flow);
+
+ private:
+  // A (node_name, output_index) pair denoting a tensor.
+  using TensorId = std::pair<string, uint32>;
+
+  // Flow file writer; defined in the .cc file.
+  class Writer;
+
+  // An operation that makes up the cell, convertible to a Myelin operation.
+  struct Operation {
+    // The TF graph node represented by this operation.
+    const tensorflow::NodeDef *node = nullptr;
+
+    // Myelin variable names of inputs to this operation.  Order matters.
+    std::vector<string> inputs;
+
+    // Number of outputs observed for this operation.  Some of the outputs in
+    // [0,|num_outputs|) might not actually be used in the cell, but we must
+    // create variables for all of them to match the expected output arity and
+    // ordering of the operation.
+    uint32 num_outputs = 0;
+  };
+
+  // Creates an empty converter.
+  MyelinCellConverter() = default;
+
+  // Implements the static Convert() method.
+  tensorflow::Status ConvertImpl(const string &component_name,
+                                 const TrainedModel &trained_model,
+                                 string *flow);
+
+  // Populates the |inputs_| and |outputs_| based on the |spec|.  On error,
+  // returns non-OK.
+  tensorflow::Status BuildInputsAndOutputs(const CellSubgraphSpec &spec);
+
+  // Walks from the |outputs_| to the |inputs_| in the |trained_model_|, adding
+  // to |operations_| along the way.  Requires that BuildInputsAndOutputs() was
+  // called.  On error, returns non-OK.
+  tensorflow::Status BuildOperations();
+
+  // Writes each section of a flow file to the |writer|.
+  tensorflow::Status WriteVariables(Writer *writer) const;
+  void WriteOperations(Writer *writer) const;
+  void WriteFunctions(Writer *writer) const;
+  void WriteConnectors(Writer *writer) const;
+  void WriteBlobs(Writer *writer) const;
+
+  // Writes a variable for the |output_index|'th output of the |node| to the
+  // |writer|.  Retrieves constant variable data from the |trained_model_| if
+  // necessary.  On error, returns non-OK.
+  tensorflow::Status WriteVariable(const tensorflow::NodeDef &node,
+                                   uint32 output_index, Writer *writer) const;
+
+  // Writes the |operation| to the |writer|.
+  void WriteOperation(const Operation &operation, Writer *writer) const;
+
+  // Returns the set of aliases associated with the |tensor_id|.
+  std::set<string> GetAliases(const TensorId &tensor_id) const;
+
+  // Parses a |tensor_name| into a |tensor_id|.  E.g.,
+  //   "foo/bar:1" => ("foo/bar", 1)
+  //   "baz"       => ("baz", 0)
+  // On error, returns non-OK.  It is an error if the |tensor_name| denotes a
+  // control dependency.
+  static tensorflow::Status ParseTensorId(const string &tensor_name,
+                                          TensorId *tensor_id);
+
+  // Returns the canonically-formatted name of the Myelin variable associated
+  // with the |tensor_id|.
+  static string AsVariableName(const TensorId &tensor_id);
+
+  // Name of the component being converted.
+  string component_name_;
+
+  // Trained model that contains the DRAGNN model.
+  const TrainedModel *trained_model_ = nullptr;
+
+  // Mapping from input tensor to logical input name.
+  std::map<TensorId, string> inputs_;
+
+  // Mapping from output tensor to logical output names.  There may be more than
+  // one name due to layer aliases (e.g., "last_layer").
+  std::map<TensorId, std::set<string>> outputs_;
+
+  // Mapping from node name to Operation.
+  std::map<string, Operation> operations_;
+};
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_

research/syntaxnet/dragnn/runtime/myelin/myelin_cell_converter_test.cc

+// Copyright 2017 Google Inc. 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 "dragnn/runtime/myelin/myelin_cell_converter.h"
+
+#include <string.h>
+#include <iostream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "dragnn/core/test/generic.h"
+#include "dragnn/runtime/alignment.h"
+#include "dragnn/runtime/myelin/myelin_spec_utils.h"
+#include "dragnn/runtime/trained_model.h"
+#include "syntaxnet/base.h"
+#include "sling/myelin/compute.h"
+#include "sling/myelin/flow.h"
+#include "sling/myelin/graph.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+// Relative path to a saved model.
+constexpr char kSavedModelDir[] = "dragnn/runtime/testdata/rnn_tagger";
+
+// Names of components in the saved model.
+const char *kComponentNames[] = {"rnn", "tagger"};
+
+// Returns a valid saved model directory.
+string GetSavedModelDir() {
+  return tensorflow::io::JoinPath(test::GetTestDataPrefix(), kSavedModelDir);
+}
+
+// Returns a string like "1048576 bytes (1.0MiB)".
+string FormatSize(int64 size) {
+  return tensorflow::strings::StrCat(
+      size, " bytes (", tensorflow::strings::HumanReadableNumBytes(size), ")");
+}
+
+// Logs the |flow|, using the |description| in the log messages.
+void DumpFlow(const sling::myelin::Flow &flow, const string &description) {
+  VLOG(1) << description << " Flow:\n" << flow.ToString();
+
+  // Log messages are truncated when they get too long.  Dump the DOT file to
+  // stdout so we get the whole thing.
+  std::cout << description << " DOT:\n"
+            << sling::myelin::FlowToDotGraph(flow, {}) << std::endl;
+}
+
+// Returns true if the |variable| is a function input or output.
+bool IsFunctionInputOrOutput(const sling::myelin::Flow::Variable &variable) {
+  // Inputs and outputs are marked with special aliases.
+  for (tensorflow::StringPiece alias : variable.aliases) {
+    if (tensorflow::str_util::StartsWith(alias, "INPUT/")) return true;
+    if (tensorflow::str_util::StartsWith(alias, "OUTPUT/")) return true;
+  }
+  return false;
+}
+
+// Returns a list of (tensor,array) pairs, one for each input and output of the
+// |flow| and |network|.  The arrays are zero-filled.
+std::vector<std::pair<sling::myelin::Tensor *, UniqueAlignedArray>>
+GetInputAndOutputTensorsAndArrays(const sling::myelin::Flow &flow,
+                                  const sling::myelin::Network &network) {
+  std::vector<std::pair<sling::myelin::Tensor *, UniqueAlignedArray>>
+      tensors_and_arrays;
+
+  for (const sling::myelin::Flow::Variable *variable : flow.vars()) {
+    // NB: Gating on variable->in || variable->out is too coarse, because that
+    // also includes constants.
+    if (!IsFunctionInputOrOutput(*variable)) continue;
+
+    sling::myelin::Tensor *tensor = network.GetParameter(variable->name);
+    CHECK(tensor != nullptr)
+        << "Failed to find tensor for variable " << variable->name;
+
+    // Currently, inputs and outputs are either int32 or float, which are the
+    // same size and have the same representation of zero.  Therefore, we can
+    // treat them the same at the byte level.
+    CHECK(tensor->type() == sling::myelin::DT_FLOAT ||
+          tensor->type() == sling::myelin::DT_INT32);
+    static_assert(sizeof(int32) == sizeof(float), "Unexpected size mismatch");
+    const int bytes = variable->shape.elements() * sizeof(float);
+
+    UniqueAlignedArray array;
+    array.Reset(bytes);
+    memset(array.view().data(), 0, bytes);  // zero for int32 or float
+
+    tensors_and_arrays.emplace_back(tensor, std::move(array));
+    VLOG(1) << "Created array of " << bytes << " bytes for variable "
+            << variable->name << " with aliases "
+            << tensorflow::str_util::Join(variable->aliases, ", ");
+  }
+
+  return tensors_and_arrays;
+}
+
+// Loads a trained model, converts it into a Flow, and then analyzes, compiles,
+// and runs the Flow.
+TEST(MyelinCellConverterTest, LoadConvertAndRun) {
+  TrainedModel trained_model;
+  TF_ASSERT_OK(trained_model.Reset(GetSavedModelDir()));
+
+  for (const string component_name : kComponentNames) {
+    LOG(INFO) << "Component: " << component_name;
+    string data;
+    TF_ASSERT_OK(
+        MyelinCellConverter::Convert(component_name, trained_model, &data));
+    LOG(INFO) << component_name << " flow size = " << FormatSize(data.size());
+
+    sling::myelin::Flow flow;
+    flow.Read(data.data(), data.size());
+
+    sling::myelin::Library library;
+    RegisterMyelinLibraries(&library);
+
+    DumpFlow(flow, tensorflow::strings::StrCat(component_name, " original"));
+    flow.Analyze(library);
+    DumpFlow(flow, tensorflow::strings::StrCat(component_name, " analyzed"));
+
+    sling::myelin::Network network;
+    ASSERT_TRUE(network.Compile(flow, library));
+
+    const sling::myelin::Cell *cell = network.GetCell(component_name);
+    ASSERT_TRUE(cell != nullptr);
+    LOG(INFO) << component_name
+              << " code size = " << FormatSize(cell->code().size());
+
+    // Create an instance and point input/output references at arrays.
+    sling::myelin::Instance instance(cell);
+    const auto tensors_and_arrays =
+        GetInputAndOutputTensorsAndArrays(flow, network);
+    for (const auto &tensor_and_array : tensors_and_arrays) {
+      instance.SetReference(tensor_and_array.first,
+                            tensor_and_array.second.view().data());
+    }
+
+    // This is just a "don't crash" test.  Myelin behavior will be exercised
+    // more thoroughly in regression tests.
+    LOG(INFO) << "Running " << component_name;
+    instance.Compute();
+    LOG(INFO) << "Successfully ran " << component_name << "!";
+  }
+}
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/myelin_dynamic_component.cc

+// Copyright 2017 Google Inc. 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 <stddef.h>
+#include <string>
+
+#include "dragnn/core/compute_session.h"
+#include "dragnn/protos/spec.pb.h"
+#include "dragnn/protos/trace.pb.h"
+#include "dragnn/runtime/component.h"
+#include "dragnn/runtime/fixed_embeddings.h"
+#include "dragnn/runtime/linked_embeddings.h"
+#include "dragnn/runtime/math/types.h"
+#include "dragnn/runtime/myelin/myelin_dynamic_component_base.h"
+#include "dragnn/runtime/network_states.h"
+#include "dragnn/runtime/session_state.h"
+#include "syntaxnet/base.h"
+#include "sling/myelin/compute.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+// A Myelin-based version of DynamicComponent.
+
+//
+// This implementation of MyelinDynamicComponentBase has the most generality
+// w.r.t. input features and links, but suffers from ComputeSession overhead.
+class MyelinDynamicComponent : public MyelinDynamicComponentBase {
+ public:
+  // Implements Component.
+  tensorflow::Status Evaluate(SessionState *session_state,
+                              ComputeSession *compute_session,
+                              ComponentTrace *component_trace) const override;
+
+ protected:
+  // Unlike other specializations, this component will only be active if the
+  // spec is explicitly modified to support Myelin (and flow resources are
+  // generated).
+  bool Supports(const ComponentSpec &spec,
+                const string &normalized_builder_name) const override {
+    return normalized_builder_name == "MyelinDynamicComponent";
+  }
+  bool PreferredTo(const Component &other) const override { return false; }
+
+ private:
+  // Forbid batches and beams.
+  static constexpr int kEvaluateNumItems = 1;
+};
+
+tensorflow::Status MyelinDynamicComponent::Evaluate(
+    SessionState *session_state, ComputeSession *compute_session,
+    ComponentTrace *component_trace) const {
+  NetworkStates &network_states = session_state->network_states;
+  FixedEmbeddings &fixed_embeddings = GetFixedEmbeddings(session_state);
+  LinkedEmbeddings &linked_embeddings = GetLinkedEmbeddings(session_state);
+
+  sling::myelin::Instance &instance = GetInstance(session_state);
+  for (size_t step_index = 0; !compute_session->IsTerminal(name());
+       ++step_index) {
+    network_states.AddStep();
+    TF_RETURN_IF_ERROR(fixed_embeddings.Reset(&fixed_embedding_manager(),
+                                              network_states, compute_session));
+    TF_RETURN_IF_ERROR(linked_embeddings.Reset(
+        &linked_embedding_manager(), network_states, compute_session));
+
+    // Bind inputs and outputs into the |instance|.
+    BindInputIds(fixed_embeddings, &instance);
+    BindInputLinks(linked_embeddings, &instance);
+    BindInputRecurrences(step_index, network_states, &instance);
+    BindOutputLayers(step_index, network_states, &instance);
+
+    // Invoke the cell in the |instance|.
+    instance.Compute();
+    MaybeTrace(step_index, &instance, component_trace);
+
+    // If the component is deterministic, take the oracle transition instead of
+    // predicting the next transition using the logits.
+    if (deterministic()) {
+      compute_session->AdvanceFromOracle(name());
+    } else {
+      // AddStep() may invalidate the logits (due to reallocation), so the layer
+      // lookup cannot be hoisted out of this loop.
+      const Vector<float> logits(
+          network_states.GetLayer(logits_handle()).row(step_index));
+      if (!compute_session->AdvanceFromPrediction(
+              name(), logits.data(), kEvaluateNumItems, logits.size())) {
+        return tensorflow::errors::Internal(
+            "Error in ComputeSession::AdvanceFromPrediction()");
+      }
+    }
+  }
+
+  return tensorflow::Status::OK();
+}
+
+DRAGNN_RUNTIME_REGISTER_COMPONENT(MyelinDynamicComponent);
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/myelin_dynamic_component_base.cc

+// Copyright 2017 Google Inc. 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 "dragnn/runtime/myelin/myelin_dynamic_component_base.h"
+
+#include <string.h>
+#include <algorithm>
+#include <set>
+
+#include "dragnn/runtime/myelin/myelin_spec_utils.h"
+#include "dragnn/runtime/transition_system_traits.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+constexpr char MyelinDynamicComponentBase::kLogitsName[];
+
+tensorflow::Status MyelinDynamicComponentBase::Validate(
+    const ComponentSpec &component_spec) {
+  if (!component_spec.attention_component().empty()) {
+    return tensorflow::errors::Unimplemented("Attention is not supported");
+  }
+
+  for (const auto &fixed_feature : component_spec.fixed_feature()) {
+    if (fixed_feature.embedding_dim() != -1) {
+      return tensorflow::errors::InvalidArgument(
+          "Myelin requires non-embedded fixed features");
+    }
+  }
+
+  for (const auto &linked_feature : component_spec.linked_feature()) {
+    if (linked_feature.embedding_dim() != -1) {
+      return tensorflow::errors::InvalidArgument(
+          "Myelin requires non-multiplied linked features");
+    }
+  }
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::LookupVector(
+    const string &name, sling::myelin::Type type, int dimension,
+    sling::myelin::Tensor **vector) const {
+  *vector = nullptr;  // so it is null if we error out
+  sling::myelin::Tensor *tensor = network_.GetParameter(name);
+  if (tensor == nullptr) {
+    return tensorflow::errors::NotFound("No Myelin tensor named '", name, "'");
+  }
+
+  if (tensor->type() != type) {
+    return tensorflow::errors::InvalidArgument(
+        "Myelin tensor has wrong type '", name, "' ", tensor->TypeString(),
+        " (expected ", sling::myelin::TypeTraits::of(type).name(), ")");
+  }
+
+  int num_nontrivial_dims = 0;
+  for (int i = 0; i < tensor->rank(); ++i) {
+    if (tensor->dim(i) > 1) ++num_nontrivial_dims;
+  }
+  if (num_nontrivial_dims > 1) {
+    return tensorflow::errors::InvalidArgument(
+        "Myelin tensor has non-vector-like shape: '", name, "' ",
+        tensor->TypeString());
+  }
+
+  // Since the |tensor| is vector-like, elements() is equivalent to the vector
+  // dimension and smooths over edges like rank=0.
+  if (dimension >= 0 && tensor->elements() != dimension) {
+    return tensorflow::errors::InvalidArgument(
+        "Myelin vector has the wrong dimension '", name, "' ",
+        tensor->TypeString(), " (expected ", dimension, ")");
+  }
+
+  if (internal::kAlignmentBytes % tensor->byte_alignment() != 0) {
+    return tensorflow::errors::FailedPrecondition(
+        "Myelin vector '", name, "' has incompatible byte alignment ",
+        tensor->byte_alignment(), " (vs ", internal::kAlignmentBytes, ")");
+  }
+
+  for (int i = 0; i < tensor->rank(); ++i) {
+    if (internal::kAlignmentBytes % tensor->minalign(i) != 0) {
+      return tensorflow::errors::FailedPrecondition(
+          "Myelin vector '", name, "' has incompatible minimum alignment ",
+          tensor->minalign(i), " for dimension ", i, " (vs ",
+          internal::kAlignmentBytes, ")");
+    }
+  }
+
+  // Success; update |vector| to non-null.
+  *vector = tensor;
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::InitializeInputIds() {
+  const int num_channels = fixed_embedding_manager_.num_channels();
+  input_ids_.resize(fixed_embedding_manager_.num_embeddings());
+  for (int channel_id = 0; channel_id < num_channels; ++channel_id) {
+    DCHECK(!fixed_embedding_manager_.is_embedded(channel_id));
+    const int channel_base = fixed_embedding_manager_.channel_base(channel_id);
+    const int channel_size = fixed_embedding_manager_.channel_size(channel_id);
+    for (int index = 0; index < channel_size; ++index) {
+      InputId &input = input_ids_[channel_base + index];
+      const string name = MakeMyelinInputFixedFeatureIdName(channel_id, index);
+      TF_RETURN_IF_ERROR(
+          LookupVector(name, sling::myelin::DT_INT32, 1, &input.id));
+      VLOG(1) << "Component '" << name_ << "' fixed channel " << channel_id
+              << " index " << index << ": Added feature ID";
+    }
+  }
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::InitializeInputLinks() {
+  const int num_channels = linked_embedding_manager_.num_channels();
+  input_links_.resize(num_channels);
+  for (int channel_id = 0; channel_id < num_channels; ++channel_id) {
+    InputLink &input = input_links_[channel_id];
+    const int dimension = linked_embedding_manager_.embedding_dim(channel_id);
+    const string activations_name =
+        MakeMyelinInputLinkedActivationVectorName(channel_id);
+    const string out_of_bounds_name =
+        MakeMyelinInputLinkedOutOfBoundsIndicatorName(channel_id);
+    TF_RETURN_IF_ERROR(LookupVector(activations_name, sling::myelin::DT_FLOAT,
+                                    dimension, &input.activations));
+    VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
+            << ": Added activations";
+
+    // Allow NOT_FOUND, for linked embedding channels that don't multiply the
+    // input activations with an embedding matrix.
+    const tensorflow::Status status = LookupVector(
+        out_of_bounds_name, sling::myelin::DT_FLOAT, 1, &input.out_of_bounds);
+    if (status.ok()) {
+      VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
+              << ": Added out-of-bounds indicator for multiplication";
+    } else if (status.code() == tensorflow::error::NOT_FOUND) {
+      VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
+              << ": No out-of-bounds indicator; not multiplied";
+    } else {
+      return status;
+    }
+  }
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::InitializeInputRecurrences(
+    const sling::myelin::Flow &flow, const NetworkStateManager &manager) {
+  for (const string &layer_name : GetRecurrentLayerNames(flow)) {
+    input_recurrences_.emplace_back();
+    InputRecurrence &input = input_recurrences_.back();
+    const string name = MakeMyelinInputRecurrentLayerName(layer_name);
+    size_t dimension = 1;
+    TF_RETURN_IF_ERROR(
+        manager.LookupLayer(name_, layer_name, &dimension, &input.handle));
+    TF_RETURN_IF_ERROR(LookupVector(name, sling::myelin::DT_FLOAT, dimension,
+                                    &input.previous_output));
+    VLOG(1) << "Component '" << name_ << "' recurrence '" << layer_name
+            << "': Added link to previous output";
+  }
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::InitializeOutputLayers(
+    const sling::myelin::Flow &flow, NetworkStateManager *manager) {
+  // Mapping from cell output tensor to layer name, for detecting layer aliases.
+  std::map<const sling::myelin::Tensor *, string> tensor_to_layer;
+  for (const string &layer_name : GetOutputLayerNames(flow)) {
+    output_layers_.emplace_back();
+    OutputLayer &output = output_layers_.back();
+    const string name = MakeMyelinOutputLayerName(layer_name);
+    TF_RETURN_IF_ERROR(
+        LookupVector(name, sling::myelin::DT_FLOAT, -1, &output.layer));
+
+    // Add a new output layer or create an alias to an existing one.
+    if (tensor_to_layer.find(output.layer) == tensor_to_layer.end()) {
+      tensor_to_layer[output.layer] = layer_name;
+      const size_t dimension = output.layer->elements();
+      TF_RETURN_IF_ERROR(
+          manager->AddLayer(layer_name, dimension, &output.handle));
+      VLOG(1) << "Component '" << name_ << "' output '" << layer_name
+              << "': Added new layer";
+    } else {
+      const string &original_name = tensor_to_layer[output.layer];
+      output_layers_.pop_back();  // not a "real" output
+      TF_RETURN_IF_ERROR(manager->AddLayerAlias(layer_name, original_name));
+      VLOG(1) << "Component '" << name_ << "' output '" << layer_name
+              << "': Alias of '" << original_name << "'";
+    }
+  }
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::InitializeConstantVectors() {
+  // Find the maximum recurrent layer dimension; the |zeros_| must be this big.
+  int max_dimension = 1;  // ensure at least one element, for |zero_|
+  for (const InputRecurrence &input : input_recurrences_) {
+    max_dimension = std::max(max_dimension, input.previous_output->elements());
+  }
+
+  // Allocate the backing array and parcel it out into sub-views.
+  const std::vector<size_t> sizes = {sizeof(float),
+                                     max_dimension * sizeof(float)};
+  array_.Reset(ComputeTotalBytesWithAlignmentPadding(sizes));
+  memset(array_.view().data(), 0, array_.view().size());  // = 0.0 for float
+  std::vector<MutableAlignedView> views;
+  TF_RETURN_IF_ERROR(array_.view().Split(sizes, &views));
+  DCHECK_EQ(views.size(), 2);
+
+  // Promote to typed vectors.
+  one_ = Vector<float>(views[0]);
+  zero_ = Vector<float>(views[1], 1);
+  zeros_ = Vector<float>(views[1]);
+  DCHECK_EQ(zero_.size(), 1);
+  DCHECK_EQ(one_.size(), 1);
+  DCHECK_EQ(zeros_.size(), max_dimension);
+
+  // All memory was already zeroed, so only |one_| needs to be initialized.
+  MutableVector<float> mutable_one(views[0]);
+  mutable_one[0] = 1.0;
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::MaybeInitializeLogits(
+    const ComponentSpec &component_spec, const NetworkStateManager &manager) {
+  // Logits are unnecessary when the component is deterministic.
+  deterministic_ = TransitionSystemTraits(component_spec).is_deterministic;
+  if (deterministic_) return tensorflow::Status::OK();
+
+  size_t dimension = 0;
+  TF_RETURN_IF_ERROR(
+      manager.LookupLayer(name_, kLogitsName, &dimension, &logits_handle_));
+
+  if (dimension != component_spec.num_actions()) {
+    return tensorflow::errors::InvalidArgument(
+        "Dimension mismatch between classification logits (", dimension,
+        ") and ComponentSpec.num_actions (", component_spec.num_actions(),
+        ") in component '", name_, "'");
+  }
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MyelinDynamicComponentBase::Initialize(
+    const ComponentSpec &component_spec, VariableStore *variable_store,
+    NetworkStateManager *network_state_manager,
+    ExtensionManager *extension_manager) {
+  name_ = component_spec.name();
+  TF_RETURN_IF_ERROR(Validate(component_spec));
+
+  const Resource *resource = nullptr;
+  TF_RETURN_IF_ERROR(LookupMyelinFlowResource(component_spec, &resource));
+  const string &flow_path = resource->part(0).file_pattern();
+
+  sling::myelin::Flow flow;
+  TF_RETURN_IF_ERROR(LoadMyelinFlow(flow_path, &flow));
+  VLOG(1) << "Original Flow for '" << name_ << "':\n" << flow.ToString();
+
+  // TODO(googleuser): Add support for optional profiling, via something like:
+  // if (...) network_.set_profiling(true)
+  // network_.Compile(flow, library);
+  // ...
+  // instance->Compute();
+  // sling::myelin::Profile profile(instance);
+  // VLOG(1) << profile.ASCIIReport();
+  RegisterMyelinLibraries(&library_);
+  flow.Analyze(library_);
+  VLOG(1) << "Analyzed Flow for '" << name_ << "':\n" << flow.ToString();
+  if (!network_.Compile(flow, library_)) {
+    return tensorflow::errors::Internal(
+        "Failed to compile Myelin network for component '", name_, "'");
+  }
+
+  cell_ = network_.GetCell(name_);
+  if (cell_ == nullptr) {
+    return tensorflow::errors::FailedPrecondition(
+        "No function named '", name_, "' in Myelin network for component '",
+        name_, "'");
+  }
+  VLOG(1) << name_ << ": " << cell_->code().size() << " bytes of Myelin code";
+
+  // Configure the inputs and outputs of the Myelin cell.  As with NetworkUnit
+  // and NetworkUnitBase, output layers and input features must be initialized
+  // in a particular order to enable recurrent inputs.  Specifically, we must
+  // populate output layers first, so they are available for recurrent access,
+  // both by the |input_recurrences_| and the |linked_embedding_manager_|.
+  TF_RETURN_IF_ERROR(InitializeOutputLayers(flow, network_state_manager));
+
+  TF_RETURN_IF_ERROR(fixed_embedding_manager_.Reset(
+      component_spec, variable_store, network_state_manager));
+  TF_RETURN_IF_ERROR(linked_embedding_manager_.Reset(
+      component_spec, variable_store, network_state_manager));
+
+  TF_RETURN_IF_ERROR(InitializeInputIds());
+  TF_RETURN_IF_ERROR(InitializeInputLinks());
+  TF_RETURN_IF_ERROR(InitializeInputRecurrences(flow, *network_state_manager));
+
+  TF_RETURN_IF_ERROR(InitializeConstantVectors());
+  TF_RETURN_IF_ERROR(
+      MaybeInitializeLogits(component_spec, *network_state_manager));
+
+  extension_manager->GetShared(&fixed_embeddings_handle_);
+  extension_manager->GetShared(&linked_embeddings_handle_);
+  extension_manager->AddLocal(&instance_handle_);
+  return tensorflow::Status::OK();
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet

research/syntaxnet/dragnn/runtime/myelin/myelin_dynamic_component_base.h

+// Copyright 2017 Google Inc. 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 DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_
+#define DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_
+
+#include <stddef.h>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "dragnn/protos/cell_trace.pb.h"
+#include "dragnn/protos/spec.pb.h"
+#include "dragnn/protos/trace.pb.h"
+#include "dragnn/runtime/alignment.h"
+#include "dragnn/runtime/component.h"
+#include "dragnn/runtime/extensions.h"
+#include "dragnn/runtime/fixed_embeddings.h"
+#include "dragnn/runtime/linked_embeddings.h"
+#include "dragnn/runtime/math/types.h"
+#include "dragnn/runtime/myelin/myelin_tracing.h"
+#include "dragnn/runtime/network_states.h"
+#include "dragnn/runtime/session_state.h"
+#include "dragnn/runtime/variable_store.h"
+#include "syntaxnet/base.h"
+#include "sling/myelin/compute.h"
+#include "sling/myelin/flow.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/dynamic_annotations.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+
+// Base class for Myelin-based versions of DynamicComponent.
+
+//
+// Roughly, this is a base class for a version of DynamicComponent where the
+// per-transition-step computation is performed by a Myelin cell instead of a
+// NetworkUnit.  This class implements Initialize() and provides methods that
+// can be useful for inference, but does not implement Evaluate().
+//
+// At initialization time, this class creates lists of configuration structs
+// that associate each input or output of the Myelin cell with an operand that
+// the DRAGNN runtime manages.  See, e.g., InputId and InitializeInputIds().
+//
+// At inference time, subclasses can bind the relevant DRAGNN runtime operands
+// to the inputs and outputs of the Myelin instance (see, e.g., BindInputIds())
+// and evaluate the Myelin cell.  Like DynamicComponent, the cell should be
+// evaluated once per transition and the results used to advance the transition
+// system state.
+//
+// Except as noted below, this is a drop-in replacement for DynamicComponent:
+// * The name of the logits layer is hard-coded (see kLogitsName).
+// * The fixed and linked channels must have embedding_dim=-1, because the fixed
+//   lookups and linked multiplications are handled within Myelin.
+//
+// The MyelinDynamicComponent subclass provides a general-purpose implementation
+// of Evaluate().  Other subclasses provide optimized implementations subject to
+// restrictions on the possible network configuration.
+class MyelinDynamicComponentBase : public Component {
+ public:
+  // Partially implements Component.
+  tensorflow::Status Initialize(const ComponentSpec &component_spec,
+                                VariableStore *variable_store,
+                                NetworkStateManager *network_state_manager,
+                                ExtensionManager *extension_manager) override;
+
+ protected:
+  // Configuration for a fixed feature ID input.
+  //
+  // TODO(googleuser): Consider making singleton inputs like the feature ID and
+  // out-of-bounds indicator into plain value inputs instead of references; it
+  // is equally fast to copy the value.
+  struct InputId {
+    // Tensor to feed with the fixed feature ID.
+    sling::myelin::Tensor *id = nullptr;
+  };
+
+  // Configuration for a linked feature embedding input.
+  struct InputLink {
+    // Tensor to feed with the linked activation vector.
+    sling::myelin::Tensor *activations = nullptr;
+
+    // Tensor to feed with the linked out-of-bounds indicator, or null if the
+    // embedding does not need to be multiplied.
+    sling::myelin::Tensor *out_of_bounds = nullptr;
+  };
+
+  // Configuration for a recurrent input.
+  struct InputRecurrence {
+    // Handle of the output layer that is recurrently fed back.
+    LayerHandle<float> handle;
+
+    // Tensor to feed with the previous output activation vector.
+    sling::myelin::Tensor *previous_output = nullptr;
+  };
+
+  // Configuration for an output layer.
+  struct OutputLayer {
+    // Handle of the output layer.
+    LayerHandle<float> handle;
+
+    // Tensor that writes to the layer.
+    sling::myelin::Tensor *layer = nullptr;
+  };
+
+  // Name of the layer containing logits.  Unlike DynamicComponent, this class
+  // does not use the NetworkUnit abstraction and assumes that the logits will
+  // be stored in this layer.
+  // TODO(googleuser): Make this configurable, if needed.  The logits layer could
+  // be given a special alias, for example.
+  static constexpr char kLogitsName[] = "logits";
+
+  // Points the cell input |tensor| in the |instance| at the |vector|.
+  template <class T>
+  static void BindInput(Vector<T> vector, sling::myelin::Tensor *tensor,
+                        sling::myelin::Instance *instance);
+
+  // Points the cell output |tensor| in the |instance| at the |vector|.
+  template <class T>
+  static void BindOutput(MutableVector<T> vector, sling::myelin::Tensor *tensor,
+                         sling::myelin::Instance *instance);
+
+  // Binds the feature IDs in the |fixed_embeddings| to the |instance| as
+  // configured by the |input_ids_|.
+  void BindInputIds(const FixedEmbeddings &fixed_embeddings,
+                    sling::myelin::Instance *instance) const;
+
+  // Binds the |embedding| and, if applicable, |is_out_of_bounds| to the
+  // |input_link| in the |instance|.
+  void BindInputLink(Vector<float> embedding, bool is_out_of_bounds,
+                     const InputLink &input_link,
+                     sling::myelin::Instance *instance) const;
+
+  // Binds the activation vectors in the |linked_embeddings| to the |instance|
+  // as configured by the |input_links_|.
+  void BindInputLinks(const LinkedEmbeddings &linked_embeddings,
+                      sling::myelin::Instance *instance) const;
+
+  // Binds the output of the step before |step_index| in the |network_states| to
+  // the |instance| as configured by the |input_recurrences_|.
+  void BindInputRecurrences(size_t step_index,
+                            const NetworkStates &network_states,
+                            sling::myelin::Instance *instance) const;
+
+  // Binds the output layers for the |step_index| in the |network_states| to the
+  // |instance| as configured by the |output_layers_|.
+  void BindOutputLayers(size_t step_index, const NetworkStates &network_states,
+                        sling::myelin::Instance *instance) const;
+
+  // Returns the reusable fixed and linked embeddings in the |session_state|.
+  FixedEmbeddings &GetFixedEmbeddings(SessionState *session_state) const;
+  LinkedEmbeddings &GetLinkedEmbeddings(SessionState *session_state) const;
+
+  // Returns the reusable Myelin instance in the |session_state|.
+  sling::myelin::Instance &GetInstance(SessionState *session_state) const;
+
+  // If |component_trace| is non-null, ensures that |step_index|+1 steps exist
+  // and traces the |instance| in the |step_index|'th step.
+  void MaybeTrace(size_t step_index, sling::myelin::Instance *instance,
+                  ComponentTrace *component_trace) const;
+
+  // Accessors.
+  const string &name() const { return name_; }
+  const FixedEmbeddingManager &fixed_embedding_manager() const {
+    return fixed_embedding_manager_;
+  }
+  const LinkedEmbeddingManager &linked_embedding_manager() const {
+    return linked_embedding_manager_;
+  }
+  const sling::myelin::Cell *cell() const { return cell_; }
+  const std::vector<InputId> &input_ids() const { return input_ids_; }
+  const std::vector<InputLink> &input_links() const { return input_links_; }
+  const std::vector<InputRecurrence> &input_recurrences() const {
+    return input_recurrences_;
+  }
+  const std::vector<OutputLayer> &output_layers() const {
+    return output_layers_;
+  }
+  bool deterministic() const { return deterministic_; }
+  LayerHandle<float> logits_handle() const { return logits_handle_; }
+
+ private:
+  // Returns non-OK if the |component_spec| specifies any unsupported settings.
+  // This includes both settings that are not yet implemented and those that are
+  // fundamentally incompatible with this class.
+  static tensorflow::Status Validate(const ComponentSpec &component_spec);
+
+  // Points the |vector| at the variable in the |network_| named |name|, which
+  // must have a vector-like shape (i.e., having at most one dimension > 1) and
+  // must match the |type|.  If the |dimension| is >= 0, then the |vector| must
+  // be the same size.  On error, returns non-OK and sets |vector| to nullptr.
+  // Returns NOT_FOUND iff the |name| does not name a variable.
+  tensorflow::Status LookupVector(const string &name, sling::myelin::Type type,
+                                  int dimension,
+                                  sling::myelin::Tensor **vector) const;
+
+  // Initializes the |input_ids_| based on the |fixed_embedding_manager_| and
+  // |network_|.  On error, returns non-OK.
+  tensorflow::Status InitializeInputIds();
+
+  // Initializes the |input_links_| based on the |linked_embedding_manager_| and
+  // |network_|.  On error, returns non-OK.
+  tensorflow::Status InitializeInputLinks();
+
+  // Initializes the |input_recurrences_| based on the |flow|, |manager|, and
+  // |network_|.  Requires that layers have been added to the |manager|.  On
+  // error, returns non-OK.
+  tensorflow::Status InitializeInputRecurrences(
+      const sling::myelin::Flow &flow, const NetworkStateManager &manager);
+
+  // Initializes the |output_layers_| based on the |flow|, |manager|, and
+  // |network_|. Adds layers to the |manager|.  On error, returns non-OK.
+  tensorflow::Status InitializeOutputLayers(const sling::myelin::Flow &flow,
+                                            NetworkStateManager *manager);
+
+  // Initializes the constant vectors (|zero_|, |one_|, and |zeros_|) and their
+  // backing |array_|.  Requires that the |input_recurrences_| are initialized.
+  tensorflow::Status InitializeConstantVectors();
+
+  // Initializes the |logits_handle_| based on the |component_spec| and
+  // |manager|, if needed.
+  tensorflow::Status MaybeInitializeLogits(const ComponentSpec &component_spec,
+                                           const NetworkStateManager &manager);
+
+  // Name of this component.
+  string name_;
+
+  // Managers for the fixed and linked embeddings used by the component.
+  FixedEmbeddingManager fixed_embedding_manager_;
+  LinkedEmbeddingManager linked_embedding_manager_;
+
+  // Fixed and linked embeddings.
+  SharedExtensionHandle<FixedEmbeddings> fixed_embeddings_handle_;
+  SharedExtensionHandle<LinkedEmbeddings> linked_embeddings_handle_;
+
+  // Library of Myelin kernels and transformations.
+  sling::myelin::Library library_;
+
+  // Myelin network that implements the cell computation.
+  sling::myelin::Network network_;
+
+  // Cell that contains the compiled code for this component.
+  const sling::myelin::Cell *cell_ = nullptr;
+
+  // List of fixed feature ID inputs, aligned with the relevant FixedEmbeddings.
+  std::vector<InputId> input_ids_;
+
+  // List of linked feature inputs, aligned with the relevant LinkedEmbeddings.
+  std::vector<InputLink> input_links_;
+
+  // List of recurrent input, not ordered.
+  std::vector<InputRecurrence> input_recurrences_;
+
+  // List of output layers, not ordered.
+  std::vector<OutputLayer> output_layers_;
+
+  // A few constant vectors and their backing array.
+  UniqueAlignedArray array_;
+  Vector<float> zero_;   // [0.0], for linked out-of-bounds indicators
+  Vector<float> one_;    // [1.0], for linked out-of-bounds indicators
+  Vector<float> zeros_;  // [0.0...0.0], for linked activation vectors
+
+  // Whether the transition system is deterministic.
+  bool deterministic_ = false;
+
+  // Handle to the classification logits.  Valid iff |deterministic_| is false.
+  LayerHandle<float> logits_handle_;
+
+  // Instance used to evaluate the network cell.  Local, since each component
+  // can have a different cell.
+  LocalExtensionHandle<sling::myelin::Instance> instance_handle_;
+};
+
+// Implementation details below.
+
+template <class T>
+void MyelinDynamicComponentBase::BindInput(Vector<T> vector,
+                                           sling::myelin::Tensor *tensor,
+                                           sling::myelin::Instance *instance) {
+  // Since Myelin only consumes non-const pointers, const_cast() is required.
+  // Myelin will not modify the contents of the |vector|, provided it is bound
+  // to a cell input.
+  DCHECK(tensor->in()) << tensor->name();
+  DCHECK(!tensor->out()) << tensor->name();
+  DCHECK_LE(tensor->elements(), vector.size()) << tensor->name();
+  instance->SetReference(
+      tensor,
+      const_cast<char *>(reinterpret_cast<const char *>(vector.data())));
+}
+
+template <class T>
+void MyelinDynamicComponentBase::BindOutput(MutableVector<T> vector,
+                                            sling::myelin::Tensor *tensor,
+                                            sling::myelin::Instance *instance) {
+  DCHECK(tensor->out()) << tensor->name();
+  DCHECK_EQ(tensor->elements(), vector.size()) << tensor->name();
+  instance->SetReference(tensor, reinterpret_cast<char *>(vector.data()));
+  TF_ANNOTATE_MEMORY_IS_INITIALIZED(vector.data(), vector.size() * sizeof(T));
+}
+
+inline void MyelinDynamicComponentBase::BindInputIds(
+    const FixedEmbeddings &fixed_embeddings,
+    sling::myelin::Instance *instance) const {
+  for (size_t i = 0; i < input_ids_.size(); ++i) {
+    BindInput(fixed_embeddings.ids(i), input_ids_[i].id, instance);
+  }
+}
+
+inline void MyelinDynamicComponentBase::BindInputLink(
+    Vector<float> embedding, bool is_out_of_bounds, const InputLink &input_link,
+    sling::myelin::Instance *instance) const {
+  BindInput(embedding, input_link.activations, instance);
+  if (input_link.out_of_bounds != nullptr) {
+    BindInput(is_out_of_bounds ? one_ : zero_, input_link.out_of_bounds,
+              instance);
+  }
+}
+
+inline void MyelinDynamicComponentBase::BindInputLinks(
+    const LinkedEmbeddings &linked_embeddings,
+    sling::myelin::Instance *instance) const {
+  for (size_t channel_id = 0; channel_id < input_links_.size(); ++channel_id) {
+    BindInputLink(linked_embeddings.embedding(channel_id),
+                  linked_embeddings.is_out_of_bounds(channel_id),
+                  input_links_[channel_id], instance);
+  }
+}
+
+inline void MyelinDynamicComponentBase::BindInputRecurrences(
+    size_t step_index, const NetworkStates &network_states,
+    sling::myelin::Instance *instance) const {
+  for (const InputRecurrence &input : input_recurrences_) {
+    if (step_index == 0) {
+      // The previous output is out-of-bounds, so feed a zero vector.  Recall
+      // that |zeros_| was constructed to be large enough for any recurrence.
+      BindInput(zeros_, input.previous_output, instance);
+    } else {
+      BindInput(Vector<float>(
+                    network_states.GetLayer(input.handle).row(step_index - 1)),
+                input.previous_output, instance);
+    }
+  }
+}
+
+inline void MyelinDynamicComponentBase::BindOutputLayers(
+    size_t step_index, const NetworkStates &network_states,
+    sling::myelin::Instance *instance) const {
+  for (const OutputLayer &output : output_layers_) {
+    BindOutput(network_states.GetLayer(output.handle).row(step_index),
+               output.layer, instance);
+  }
+}
+
+inline FixedEmbeddings &MyelinDynamicComponentBase::GetFixedEmbeddings(
+    SessionState *session_state) const {
+  return session_state->extensions.Get(fixed_embeddings_handle_);
+}
+
+inline LinkedEmbeddings &MyelinDynamicComponentBase::GetLinkedEmbeddings(
+    SessionState *session_state) const {
+  return session_state->extensions.Get(linked_embeddings_handle_);
+}
+
+inline sling::myelin::Instance &MyelinDynamicComponentBase::GetInstance(
+    SessionState *session_state) const {
+  return session_state->extensions.Get(instance_handle_, cell_);
+}
+
+inline void MyelinDynamicComponentBase::MaybeTrace(
+    size_t step_index, sling::myelin::Instance *instance,
+    ComponentTrace *component_trace) const {
+  if (component_trace == nullptr) return;
+  while (component_trace->step_trace_size() <= step_index) {
+    component_trace->add_step_trace();
+  }
+  TraceMyelinInstance(instance,
+                      component_trace->mutable_step_trace(step_index)
+                          ->AddExtension(CellTrace::step_trace_extension));
+}
+
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet
+
+#endif  // DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_

research/syntaxnet/dragnn/runtime/myelin/myelin_dynamic_component_test.cc

+// Copyright 2017 Google Inc. 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 <string>
+
+#include "dragnn/core/test/generic.h"
+#include "dragnn/protos/cell_trace.pb.h"
+#include "dragnn/protos/spec.pb.h"
+#include "dragnn/protos/trace.pb.h"
+#include "dragnn/runtime/component.h"
+#include "dragnn/runtime/extensions.h"
+#include "dragnn/runtime/myelin/myelin_spec_utils.h"
+#include "dragnn/runtime/test/network_test_base.h"
+#include "syntaxnet/base.h"
+#include "sling/file/file.h"
+#include "sling/myelin/flow.h"
+#include <gmock/gmock.h>
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace syntaxnet {
+namespace dragnn {
+namespace runtime {
+namespace {
+
+using ::testing::_;
+using ::testing::InSequence;
+using ::testing::Invoke;
+using ::testing::Return;
+
+constexpr int kFlowVersion = 4;
+constexpr int kVocabularySize = 123;
+constexpr int kLogitsDim = 11;
+constexpr int kNumSteps = 50;
+
+class MyelinDynamicComponentTest : public NetworkTestBase {
+ protected:
+  // Options for building a Flow file for tests.  By default, this specifies a
+  // working Flow file, but settings can be perturbed to trigger errors.
+  struct FlowFileOptions {
+    FlowFileOptions() = default;
+
+    // Name of the function to create.
+    string function_name = kTestComponentName;
+
+    // Dimension of the classification logits.
+    int logits_dim = kLogitsDim;
+
+    // Name of the variable containing the classification logits.
+    string logits_name = "logits";
+
+    // Type of the feature ID input.
+    sling::myelin::Type id_type = sling::myelin::DT_INT32;
+
+    // Dimension of the feature ID input.
+    int id_dim = 1;
+  };
+
+  // Builds and writes a simple Flow file.  By default it produces a valid Flow,
+  // but arguments can be overridden for error testing.  Returns the path to the
+  // Flow file.
+  static string WriteFlowFile() { return WriteFlowFile(FlowFileOptions()); }
+  static string WriteFlowFile(const FlowFileOptions &options) {
+    sling::myelin::Flow flow;
+
+    // A fixed feature ID input.
+    sling::myelin::Flow::Variable *id =
+        flow.AddVariable("id", options.id_type, {options.id_dim});
+    id->ref = true;
+    id->aliases.push_back(MakeMyelinInputFixedFeatureIdName(0, 0));
+
+    // An embedding matrix constant.  Each embedding is filled with its index.
+    sling::myelin::Flow::Variable *embeddings =
+        flow.AddVariable("embeddings", sling::myelin::DT_FLOAT,
+                         {kVocabularySize, options.logits_dim});
+    std::vector<float> data(kVocabularySize * options.logits_dim);
+    for (int row = 0; row < kVocabularySize; ++row) {
+      for (int column = 0; column < options.logits_dim; ++column) {
+        data[row * options.logits_dim + column] = row;
+      }
+    }
+    embeddings->SetData(data.data(), data.size() * sizeof(float));
+
+    // The retrieved embedding row, as logits.
+    sling::myelin::Flow::Variable *logits =
+        flow.AddVariable(options.logits_name, sling::myelin::DT_FLOAT,
+                         {options.id_dim, options.logits_dim});
+    logits->ref = true;
+    logits->aliases.push_back(MakeMyelinOutputLayerName(options.logits_name));
+
+    // A Gather op that looks up the |id| in the |embeddings|, and returns the
+    // result in the |logits|.
+    flow.AddOperation(flow.AddFunction(options.function_name), "gather",
+                      "Gather", {embeddings, id}, {logits});
+
+    const string flow_path =
+        tensorflow::io::JoinPath(tensorflow::testing::TmpDir(), "foo.flow");
+    sling::File::Init();
+    flow.Save(flow_path, kFlowVersion);
+    return flow_path;
+  }
+
+  // Creates a component, initializes it based on the |component_spec_text| and
+  // |flow_path|, and evaluates it.  The |component_trace| is overwritten with
+  // traces, if non-null.  On error, returns non-OK.
+  tensorflow::Status Run(const string &component_spec_text = "",
+                         const string &flow_path = WriteFlowFile(),
+                         ComponentTrace *component_trace = nullptr) {
+    ComponentSpec component_spec;
+    CHECK(TextFormat::ParseFromString(component_spec_text, &component_spec));
+    if (!component_spec.has_num_actions()) {
+      component_spec.set_num_actions(kLogitsDim);
+    }
+    component_spec.set_name(kTestComponentName);
+
+    auto *fixed_feature = component_spec.add_fixed_feature();
+    fixed_feature->set_embedding_dim(-1);
+    fixed_feature->set_size(1);
+
+    TF_RETURN_IF_ERROR(AddMyelinFlowResource(flow_path, &component_spec));
+
+    AddComponent(kTestComponentName);
+    TF_RETURN_IF_ERROR(
+        Component::CreateOrError("MyelinDynamicComponent", &component_));
+    TF_RETURN_IF_ERROR(component_->Initialize(component_spec, &variable_store_,
+                                              &network_state_manager_,
+                                              &extension_manager_));
+
+    network_states_.Reset(&network_state_manager_);
+    StartComponent(0);  // MyelinDynamicComponent will add steps
+    session_state_.extensions.Reset(&extension_manager_);
+
+    TF_RETURN_IF_ERROR(component_->Evaluate(&session_state_, &compute_session_,
+                                            component_trace));
+    return tensorflow::Status::OK();
+  }
+
+  std::unique_ptr<Component> component_;
+};
+
+// Tests that MyelinDynamicComponent fails if the spec uses attention.
+TEST_F(MyelinDynamicComponentTest, UnsupportedAttention) {
+  EXPECT_THAT(Run("attention_component:'foo'"),
+              test::IsErrorWithSubstr("Attention is not supported"));
+}
+
+// Tests that MyelinDynamicComponent fails if the spec has embedded fixed
+// features.
+TEST_F(MyelinDynamicComponentTest, InvalidFixedFeatureIsEmbedded) {
+  EXPECT_THAT(
+      Run("fixed_feature { embedding_dim:1 }"),
+      test::IsErrorWithSubstr("Myelin requires non-embedded fixed features"));
+}
+
+// Tests that MyelinDynamicComponent fails if the ComponentSpec has a fixed
+// feature that does not appear in the Flow.
+TEST_F(MyelinDynamicComponentTest, InvalidFixedFeatureNotInFlow) {
+  EXPECT_THAT(Run("fixed_feature { embedding_dim:-1 size:1 }"),
+              test::IsErrorWithSubstr(tensorflow::strings::StrCat(
+                  "No Myelin tensor named '",
+                  MakeMyelinInputFixedFeatureIdName(1, 0), "'")));
+}
+
+// Tests that MyelinDynamicComponent fails if the spec has multipled linked
+// features.
+TEST_F(MyelinDynamicComponentTest, InvalidLinkedFeatureIsMultiplied) {
+  EXPECT_THAT(Run("linked_feature { embedding_dim:1 }"),
+              test::IsErrorWithSubstr(
+                  "Myelin requires non-multiplied linked features"));
+}
+
+// Tests that MyelinDynamicComponent fails if the ComponentSpec has a linked
+// feature that does not appear in the Flow.
+TEST_F(MyelinDynamicComponentTest, InvalidLinkedFeatureNotInFlow) {
+  const string kSpec = tensorflow::strings::StrCat(
+      "linked_feature { source_component:'", kTestComponentName,
+      "' source_layer:'logits' embedding_dim:-1 size:1 }");
+
+  EXPECT_THAT(Run(kSpec),
+              test::IsErrorWithSubstr(tensorflow::strings::StrCat(
+                  "No Myelin tensor named '",
+                  MakeMyelinInputLinkedActivationVectorName(0), "'")));
+}
+
+// Tests that MyelinDynamicComponent fails if the Flow file does not exist.
+TEST_F(MyelinDynamicComponentTest, InvalidFlowFilePath) {
+  EXPECT_THAT(Run("", "/invalid/path"),
+              test::IsErrorWithSubstr("Failed to load Myelin Flow"));
+}
+
+// Tests that MyelinDynamicComponent fails if the function in the Flow file has
+// the wrong name.
+TEST_F(MyelinDynamicComponentTest, WrongFunctionName) {
+  FlowFileOptions options;
+  options.function_name = "wrong_function";
+
+  EXPECT_THAT(
+      Run("", WriteFlowFile(options)),
+      test::IsErrorWithSubstr(tensorflow::strings::StrCat(
+          "No function named '", kTestComponentName, "' in Myelin network")));
+}
+
+// Tests that MyelinDynamicComponent fails if the logits dimension does not
+// match ComponentSpec.num_actions.
+TEST_F(MyelinDynamicComponentTest, WrongLogitsDimension) {
+  FlowFileOptions options;
+  options.logits_dim = kLogitsDim + 1;
+
+  EXPECT_THAT(Run("", WriteFlowFile(options)),
+              test::IsErrorWithSubstr(
+                  "Dimension mismatch between classification logits"));
+}
+
+// Tests that MyelinDynamicComponent fails if there is no "logits" layer.
+TEST_F(MyelinDynamicComponentTest, WrongLogitsName) {
+  FlowFileOptions options;
+  options.logits_name = "not_logits";
+
+  EXPECT_THAT(Run("", WriteFlowFile(options)),
+              test::IsErrorWithSubstr("Unknown layer 'logits'"));
+}
+
+// Tests that MyelinDynamicComponent fails to compile if one of the Myelin
+// tensors has the wrong type.
+TEST_F(MyelinDynamicComponentTest, FailToCompile) {
+  FlowFileOptions options;
+  options.id_type = sling::myelin::DT_FLOAT;
+
+  EXPECT_THAT(Run("", WriteFlowFile(options)),
+              test::IsErrorWithSubstr("Failed to compile Myelin network"));
+}
+
+// Tests that MyelinDynamicComponent fails if one of the Myelin tensors is not
+// vector-like.
+TEST_F(MyelinDynamicComponentTest, NotVectorLike) {
+  FlowFileOptions options;
+  options.id_dim = 2;
+
+  EXPECT_THAT(
+      Run("", WriteFlowFile(options)),
+      test::IsErrorWithSubstr("Myelin tensor has non-vector-like shape"));
+}
+
+// Tests that MyelinDynamicComponent fails if AdvanceFromPrediction() fails.
+TEST_F(MyelinDynamicComponentTest, FailToAdvanceFromPrediction) {
+  EXPECT_CALL(compute_session_, IsTerminal(_)).WillRepeatedly(Return(false));
+  EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
+      .WillOnce(Return(false));
+  EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
+      .WillOnce(Invoke(ExtractFeatures(0, {{10, 1.0}})));
+
+  EXPECT_THAT(Run(), test::IsErrorWithSubstr(
+                         "Error in ComputeSession::AdvanceFromPrediction()"));
+}
+
+// Tests that MyelinDynamicComponent can run a simple non-deterministic Flow.
+TEST_F(MyelinDynamicComponentTest, SimpleNonDeterministicFlow) {
+  SetupTransitionLoop(kNumSteps);
+  EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
+      .Times(kNumSteps)
+      .WillRepeatedly(Return(true));
+
+  {  // Extract a sequence of feature IDs equal to 2 * step_index.
+    ASSERT_LE(2 * kNumSteps, kVocabularySize);
+    InSequence scoped;
+    for (int step_index = 0; step_index < kNumSteps; ++step_index) {
+      EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
+          .WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
+    }
+  }
+
+  TF_ASSERT_OK(Run());
+
+  const Matrix<float> logits(GetLayer(kTestComponentName, "logits"));
+  ASSERT_EQ(logits.num_rows(), kNumSteps);
+  ASSERT_EQ(logits.num_columns(), kLogitsDim);
+
+  // Since each row of the embedding matrix is filled with its index, the logits
+  // should be equal to the feature IDs.
+  for (int step_index = 0; step_index < kNumSteps; ++step_index) {
+    ExpectVector(logits.row(step_index), kLogitsDim, 2 * step_index);
+  }
+}
+
+// Tests that MyelinDynamicComponent can run a simple deterministic Flow.
+TEST_F(MyelinDynamicComponentTest, SimpleDeterministicFlow) {
+  SetupTransitionLoop(kNumSteps);
+  EXPECT_CALL(compute_session_, AdvanceFromOracle(kTestComponentName))
+      .Times(kNumSteps);
+
+  {  // Extract a sequence of feature IDs equal to 2 * step_index.
+    ASSERT_LE(2 * kNumSteps, kVocabularySize);
+    InSequence scoped;
+    for (int step_index = 0; step_index < kNumSteps; ++step_index) {
+      EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
+          .WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
+    }
+  }
+
+  FlowFileOptions options;
+  options.logits_dim = 1;
+  TF_ASSERT_OK(Run("num_actions:1", WriteFlowFile(options)));
+}
+
+// Tests that MyelinDynamicComponent can run a simple Flow with tracing enabled.
+TEST_F(MyelinDynamicComponentTest, SimpleFlowWithTracing) {
+  SetupTransitionLoop(kNumSteps);
+  EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
+      .Times(kNumSteps)
+      .WillRepeatedly(Return(true));
+
+  { // Extract a sequence of feature IDs equal to 2 * step_index.
+    ASSERT_LE(2 * kNumSteps, kVocabularySize);
+    InSequence scoped;
+    for (int step_index = 0; step_index < kNumSteps; ++step_index) {
+      EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
+          .WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
+    }
+  }
+
+  ComponentTrace component_trace;
+  TF_ASSERT_OK(Run("", WriteFlowFile(), &component_trace));
+
+  // Each step trace should have a cell trace from the Myelin instance.
+  ASSERT_EQ(component_trace.step_trace_size(), kNumSteps);
+  for (const ComponentStepTrace &step_trace : component_trace.step_trace()) {
+    ASSERT_EQ(step_trace.ExtensionSize(CellTrace::step_trace_extension), 1);
+    const CellTrace &cell_trace =
+        step_trace.GetExtension(CellTrace::step_trace_extension, 0);
+    EXPECT_EQ(cell_trace.name(), kTestComponentName);
+  }
+}
+
+}  // namespace
+}  // namespace runtime
+}  // namespace dragnn
+}  // namespace syntaxnet