1 files changed, 82 insertions, 0 deletions
diff --git a/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/SoftMax.h b/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/SoftMax.h
new file mode 100644
index 0000000000..223ae28ffd
--- /dev/null
+++ b/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/SoftMax.h
@@ -0,0 +1,82 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#ifndef EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
+#define EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
+
+namespace Eigen {
+
+/** SoftMax
+  * \ingroup CXX11_NeuralNetworks_Module
+  *
+  * \brief Applies a softmax
+  *
+  * The input parameter is expected to be a col-major tensor with a rank of 2 (depth and other).
+  *
+  * The result can be assigned to a tensor of rank and dimensions equal to that of the input. The result will be laid out in col-major order.
+  *
+*/
+
+namespace {
+struct SoftmaxOp {
+  SoftmaxOp(const float beta) : beta_(beta) { }
+
+  template <typename Input>
+  typename Input::Dimensions dimensions(const Input& input) const {
+    return input.dimensions();
+  }
+
+  template <typename Input, typename Output, typename Device>
+  void eval(const Input& input, Output& output, const Device& device) const
+  {
+#if !defined(EIGEN_HAS_INDEX_LIST)
+    // nvcc doesn't support cxx11
+    Eigen::array<typename internal::traits<Input>::Index, 1> depth_dim;
+    depth_dim[0] = 0;
+    Eigen::array<typename internal::traits<Input>::Index, 2> bcast;
+    bcast[0] = dimensions(input)[0];
+    bcast[1] = 1;
+    DSizes<typename internal::traits<Input>::Index, 2> dims2d;
+    dims2d[0] = 1;
+    dims2d[1] = dimensions(input)[1];
+#else
+    // Take advantage of cxx11 to give the compiler information it can use to
+    // optimize the code.
+    Eigen::IndexList<Eigen::type2index<0>> depth_dim;
+    Eigen::IndexList<int, Eigen::type2index<1>> bcast;
+    bcast.set(0, dimensions(input)[0]);
+    Eigen::IndexList<Eigen::type2index<1>, typename internal::traits<Input>::Index> dims2d;
+    dims2d.set(1, dimensions(input)[1]);
+#endif
+
+    output.device(device) = ((input - input.maximum(depth_dim).eval().reshape(dims2d).broadcast(bcast)) * beta_).exp();
+    output.device(device) = output / (output.sum(depth_dim).eval().reshape(dims2d).broadcast(bcast));
+  }
+
+ private:
+  const float beta_;
+};
+}
+
+
+template <typename Input>
+EIGEN_ALWAYS_INLINE
+static const TensorCustomUnaryOp<const SoftmaxOp, const Input>
+SoftMax(const Input& input, const float beta)
+{
+  EIGEN_STATIC_ASSERT(internal::traits<Input>::Layout == ColMajor, YOU_MADE_A_PROGRAMMING_MISTAKE);
+  EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions == 2, YOU_MADE_A_PROGRAMMING_MISTAKE);
+
+  const SoftmaxOp op(beta);
+  return input.customOp(op);
+}
+
+
+} // end namespace Eigen
+
+#endif // EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H