def call(self, inputs, states):
   prev_output = states[0]
   h = special_math_ops.einsum('bij,ijkl->bkl', inputs, self.kernel)
   h += array_ops.expand_dims(self.bias, axis=0)
   output = h + special_math_ops.einsum('bij,ijkl->bkl', prev_output,
                                        self.recurring_kernel)
   return output, [output]

  def test_input_is_placeholder(self):
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(1, None))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, 1))
      out = special_math_ops.einsum('ij,jk->ik', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [[2], [1], [1]],
        }
        np.testing.assert_almost_equal(
            [[7]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, 3))
      m1 = array_ops.placeholder(dtypes.int32, shape=(3,))
      out = special_math_ops.einsum('ij,j->i', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [2, 1, 1],
        }
        np.testing.assert_almost_equal([7], sess.run(out, feed_dict=feed_dict))

    # Tests for placeholders which have two or more None values
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      m1 = array_ops.placeholder(dtypes.int32, shape=(2, 1))
      out = special_math_ops.einsum('ijk,kl->ijl', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[[1,2]]],
            m1: [[3], [2]],
        }
        np.testing.assert_almost_equal(
            [[[7]]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(2, 1))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      out = special_math_ops.einsum('kl,ijk->ijl', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[3], [2]],
            m1: [[[1,2]]],
        }
        np.testing.assert_almost_equal(
            [[[7]]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      m1 = array_ops.placeholder(dtypes.int32, shape=(2,))
      out = special_math_ops.einsum('ijk,k->ij', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[[1, 2]]],
            m1: [3, 2],
        }
        np.testing.assert_almost_equal(
            [[7]], sess.run(out, feed_dict=feed_dict))

 def _convdiag_sum_of_squares(self, patches, outputs_grad):
   # This computes the sum of the squares of the per-training-case "gradients".
   # It does this simply by computing a giant tensor containing all of these,
   # doing an entry-wise square, and them summing along the batch dimension.
   case_wise_gradients = special_math_ops.einsum("bijk,bijl->bkl", patches,
                                                 outputs_grad)
   return math_ops.reduce_sum(math_ops.square(case_wise_gradients), axis=0)

 def test_invalid_keyword_arguments(self):
   m0 = array_ops.placeholder(dtypes.int32, shape=(1, None))
   m1 = array_ops.placeholder(dtypes.int32, shape=(None, 1))
   with self.assertRaisesRegexp(TypeError,
       'invalid keyword arguments for this function: invalid1, invalid2'):
     _ = special_math_ops.einsum('ij,jk->ik', m0, m1, name="name",
                                 invalid1="value1", invalid2="value2")

  def run_test(self, axes, expanded_axes=None):
    expanded_axes = expanded_axes if expanded_axes is not None else axes
    all_axes = {ax: np.random.randint(4, 12)
                for ax in expanded_axes if ax.isalpha()}

    input_vals = []
    input_axes, _, _ = axes.partition('->')

    for idx in input_axes.split(','):
      shape = [all_axes[ax] for ax in idx if ax.isalpha()]
      input_vals.append(np.random.random(shape))

    input_tensors = [constant_op.constant(val) for val in input_vals]
    output_tensor = special_math_ops.einsum(axes, *input_tensors)

    with self.session(use_gpu=True):
      output_value = self.evaluate(output_tensor)

    correct_value = 0
    if axes == 'ijji':
      output = math_ops.trace(*input_tensors)
      correct_value = self.evaluate(output)
    else:
      correct_value = np.einsum(axes, *input_vals)
    err = np.abs(correct_value - output_value).max()
    self.assertLess(err, 1e-8)

 def test_dim_mismatch(self):
   for axes, input_shapes in self.dim_mismatch_cases:
     inputs = [
         array_ops.placeholder(dtypes.float32, shape=shape)
         for shape in input_shapes
     ]
     with self.assertRaises(ValueError):
       _ = special_math_ops.einsum(axes, *inputs)

 def test_invalid(self):
   for axes in self.invalid_cases:
     inputs = [
         array_ops.placeholder(dtypes.float32, shape=(3, 4)),
         array_ops.placeholder(dtypes.float32, shape=(3, 4)),
     ]
     with self.assertRaises(ValueError):
       _ = special_math_ops.einsum(axes, *inputs)

  def test_input_is_placeholder(self):
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(1, None))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, 1))
      out = special_math_ops.einsum('ij,jk->ik', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [[2], [1], [1]],
        }
        np.testing.assert_almost_equal(
            [[7]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, 3))
      m1 = array_ops.placeholder(dtypes.int32, shape=(3,))
      out = special_math_ops.einsum('ij,j->i', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [2, 1, 1],
        }
        np.testing.assert_almost_equal([7], sess.run(out, feed_dict=feed_dict))

  def call(self, inputs, states):
    # inputs should be in [(batch, input_1), (batch, input_2, input_3)]
    # state should be in shape [(batch, unit_1), (batch, unit_2, unit_3)]
    flatten_inputs = nest.flatten(inputs)
    s1, s2 = states

    output_1 = math_ops.matmul(flatten_inputs[0], self.kernel_1)
    output_2_3 = special_math_ops.einsum('bij,ijkl->bkl', flatten_inputs[1],
                                         self.kernel_2_3)
    state_1 = s1 + output_1
    state_2_3 = s2 + output_2_3

    output = [output_1, output_2_3]
    new_states = NestedState(s1=state_1, s2=state_2_3)

    return output, new_states

  def run_test(self, axes):
    all_axes = {ax: np.random.randint(4, 12) for ax in axes if ax.isalpha()}

    input_vals = []
    input_axes, _, _ = axes.partition('->')

    for idx in input_axes.split(','):
      shape = [all_axes[ax] for ax in idx]
      input_vals.append(np.random.random(shape))

    input_tensors = [constant_op.constant(val) for val in input_vals]
    output_tensor = special_math_ops.einsum(axes, *input_tensors)

    with self.test_session(use_gpu=True):
      output_value = output_tensor.eval()

    correct_value = np.einsum(axes, *input_vals)

    err = np.abs(correct_value - output_value).max()
    print(axes, err)
    assert err < 1e-8

 def test_repeated_axis_single_input(self):
   x = array_ops.placeholder(dtypes.float32, shape=[2, 2])
   with self.assertRaises(ValueError):
     _ = special_math_ops.einsum('ii->', x)

 def test_ellipses_with_unknown_input_dim(self):
   with ops.Graph().as_default():
     m0 = array_ops.placeholder(dtypes.float32)
     m1 = array_ops.placeholder_with_default([[3, 2]], shape=(None, 2))
     with self.assertRaises(ValueError):
       _ = special_math_ops.einsum('...jkl,...j->...kl', m0, m1)

 def test_multiple_ellipses(self):
   m0 = array_ops.placeholder_with_default([[[[1, 2]], [[2, 1]]]],
                                           shape=(None, 2, None, 2))
   m1 = array_ops.placeholder_with_default([[3, 2]], shape=(None, 2))
   out = special_math_ops.einsum('...jkl,...j->...kl', m0, m1)
   self.assertAllClose([[[7, 8]]], self.evaluate(out))