spark-nlp 2.6.3rc1__py2.py3-none-any.whl → 6.2.1__py2.py3-none-any.whl

sparknlp/training/_tf_graph_builders/tf2contrib/gru_ops.py

@@ -0,0 +1,235 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Python wrapper for the Block GRU Op."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.keras.engine import input_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_rnn_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import rnn_cell_impl
+from tensorflow.python.util.deprecation import deprecated_args
+
+LayerRNNCell = rnn_cell_impl.LayerRNNCell  # pylint: disable=invalid-name
+
+
+@ops.RegisterGradient("GRUBlockCell")
+def _GRUBlockCellGrad(op, *grad):
+    r"""Gradient for GRUBlockCell.
+
+    Args:
+      op: Op for which the gradient is defined.
+      *grad: Gradients of the optimization function wrt output
+        for the Op.
+
+    Returns:
+      d_x: Gradients wrt to x
+      d_h: Gradients wrt to h
+      d_w_ru: Gradients wrt to w_ru
+      d_w_c: Gradients wrt to w_c
+      d_b_ru: Gradients wrt to b_ru
+      d_b_c: Gradients wrt to b_c
+
+    Mathematics behind the Gradients below:
+    ```
+    d_c_bar = d_h \circ (1-u) \circ (1-c \circ c)
+    d_u_bar = d_h \circ (h-c) \circ u \circ (1-u)
+
+    d_r_bar_u_bar = [d_r_bar d_u_bar]
+
+    [d_x_component_1 d_h_prev_component_1] = d_r_bar_u_bar * w_ru^T
+
+    [d_x_component_2 d_h_prevr] = d_c_bar * w_c^T
+
+    d_x = d_x_component_1 + d_x_component_2
+
+    d_h_prev = d_h_prev_component_1 + d_h_prevr \circ r + u
+    ```
+    Below calculation is performed in the python wrapper for the Gradients
+    (not in the gradient kernel.)
+    ```
+    d_w_ru = x_h_prevr^T * d_c_bar
+
+    d_w_c = x_h_prev^T * d_r_bar_u_bar
+
+    d_b_ru = sum of d_r_bar_u_bar along axis = 0
+
+    d_b_c = sum of d_c_bar along axis = 0
+    ```
+    """
+    x, h_prev, w_ru, w_c, b_ru, b_c = op.inputs
+    r, u, c, _ = op.outputs
+    _, _, _, d_h = grad
+
+    d_x, d_h_prev, d_c_bar, d_r_bar_u_bar = gen_rnn_ops.gru_block_cell_grad(
+        x, h_prev, w_ru, w_c, b_ru, b_c, r, u, c, d_h)
+
+    x_h_prev = array_ops.concat([x, h_prev], 1)
+    d_w_ru = math_ops.matmul(x_h_prev, d_r_bar_u_bar, transpose_a=True)
+    d_b_ru = nn_ops.bias_add_grad(d_r_bar_u_bar)
+
+    x_h_prevr = array_ops.concat([x, h_prev * r], 1)
+    d_w_c = math_ops.matmul(x_h_prevr, d_c_bar, transpose_a=True)
+    d_b_c = nn_ops.bias_add_grad(d_c_bar)
+
+    return d_x, d_h_prev, d_w_ru, d_w_c, d_b_ru, d_b_c
+
+
+class GRUBlockCell(LayerRNNCell):
+    r"""Block GRU cell implementation.
+
+    Deprecated: use GRUBlockCellV2 instead.
+
+    The implementation is based on:  http://arxiv.org/abs/1406.1078
+    Computes the GRU cell forward propagation for 1 time step.
+
+    This kernel op implements the following mathematical equations:
+
+    Biases are initialized with:
+
+    * `b_ru` - constant_initializer(1.0)
+    * `b_c` - constant_initializer(0.0)
+
+    ```
+    x_h_prev = [x, h_prev]
+
+    [r_bar u_bar] = x_h_prev * w_ru + b_ru
+
+    r = sigmoid(r_bar)
+    u = sigmoid(u_bar)
+
+    h_prevr = h_prev \circ r
+
+    x_h_prevr = [x h_prevr]
+
+    c_bar = x_h_prevr * w_c + b_c
+    c = tanh(c_bar)
+
+    h = (1-u) \circ c + u \circ h_prev
+    ```
+
+    """
+
+    @deprecated_args(None, "cell_size is deprecated, use num_units instead",
+                     "cell_size")
+    def __init__(self,
+                 num_units=None,
+                 cell_size=None,
+                 reuse=None,
+                 name="gru_cell"):
+        """Initialize the Block GRU cell.
+
+        Args:
+          num_units: int, The number of units in the GRU cell.
+          cell_size: int, The old (deprecated) name for `num_units`.
+          reuse: (optional) boolean describing whether to reuse variables in an
+            existing scope.  If not `True`, and the existing scope already has the
+            given variables, an error is raised.
+          name: String, the name of the layer. Layers with the same name will
+            share weights, but to avoid mistakes we require reuse=True in such
+            cases.  By default this is "lstm_cell", for variable-name compatibility
+            with `tf.compat.v1.nn.rnn_cell.GRUCell`.
+
+        Raises:
+          ValueError: if both cell_size and num_units are not None;
+            or both are None.
+        """
+        super(GRUBlockCell, self).__init__(_reuse=reuse, name=name)
+        if (cell_size is None) == (num_units is None):
+            raise ValueError(
+                "Exactly one of num_units or cell_size must be provided.")
+        if num_units is None:
+            num_units = cell_size
+        self._cell_size = num_units
+        # Inputs must be 2-dimensional.
+        self.input_spec = input_spec.InputSpec(ndim=2)
+
+    @property
+    def state_size(self):
+        return self._cell_size
+
+    @property
+    def output_size(self):
+        return self._cell_size
+
+    def build(self, input_shape):
+        # Check if the input size exist.
+        input_size = tensor_shape.dimension_value(input_shape[1])
+        if input_size is None:
+            raise ValueError("Expecting input_size to be set.")
+
+        self._gate_kernel = self.add_variable(
+            "w_ru", [input_size + self._cell_size, self._cell_size * 2])
+        self._gate_bias = self.add_variable(
+            "b_ru", [self._cell_size * 2],
+            initializer=init_ops.constant_initializer(1.0))
+        self._candidate_kernel = self.add_variable(
+            "w_c", [input_size + self._cell_size, self._cell_size])
+        self._candidate_bias = self.add_variable(
+            "b_c", [self._cell_size],
+            initializer=init_ops.constant_initializer(0.0))
+
+        self.built = True
+
+    def call(self, inputs, h_prev):
+        """GRU cell."""
+        # Check cell_size == state_size from h_prev.
+        cell_size = h_prev.get_shape().with_rank(2)[1]
+        if cell_size != self._cell_size:
+            raise ValueError("Shape of h_prev[1] incorrect: cell_size %i vs %s" %
+                             (self._cell_size, cell_size))
+
+        _gru_block_cell = gen_rnn_ops.gru_block_cell
+        _, _, _, new_h = _gru_block_cell(
+            x=inputs,
+            h_prev=h_prev,
+            w_ru=self._gate_kernel,
+            w_c=self._candidate_kernel,
+            b_ru=self._gate_bias,
+            b_c=self._candidate_bias)
+
+        return new_h, new_h
+
+
+class GRUBlockCellV2(GRUBlockCell):
+    """Temporary GRUBlockCell impl with a different variable naming scheme.
+
+    Only differs from GRUBlockCell by variable names.
+    """
+
+    def build(self, input_shape):
+        """GRU cell."""
+        input_size = tensor_shape.dimension_value(input_shape[1])
+        if input_size is None:
+            raise ValueError("Expecting input_size to be set.")
+
+        self._gate_kernel = self.add_variable(
+            "gates/kernel", [input_size + self._cell_size, self._cell_size * 2])
+        self._gate_bias = self.add_variable(
+            "gates/bias", [self._cell_size * 2],
+            initializer=init_ops.constant_initializer(1.0))
+        self._candidate_kernel = self.add_variable(
+            "candidate/kernel", [input_size + self._cell_size, self._cell_size])
+        self._candidate_bias = self.add_variable(
+            "candidate/bias", [self._cell_size],
+            initializer=init_ops.constant_initializer(0.0))
+
+        self.built = True