message: Add symbolic gradient for LogSoftmax. The symbolic gradient for Softmax is also
  simplified.
diff: @@ -37,18 +37,41 @@ Status SoftmaxGrad(const AttrSlice& attrs, FunctionDef* g) {
  {
  {{"softmax"}, "Softmax", {"x"}, {{"T", "$T"}}},
  {{"n0"}, "Mul", {"grad_softmax", "softmax"}, {{"T", "$T"}}},
  - FDH::Const<int32>("indices", {1}),
  - {{"n1"}, "Sum", {"n0", "indices"}, {{"T", "$T"}}},
  - FDH::Const<int32>("newshape", {-1, 1}),
  - {{"n2"}, "Reshape", {"n1", "newshape"}, {{"T", "$T"}}},
  - {{"n3"}, "Sub", {"grad_softmax", "n2"}, {{"T", "$T"}}},
  - {{"grad_x"}, "Mul", {"n3", "softmax"}, {{"T", "$T"}}}
  + FDH::Const<int32>("indices", {-1}),
  + {{"n1"}, "Sum", {"n0", "indices"}, {{"keep_dims", true}, {"T", "$T"}}},
  + {{"n2"}, "Sub", {"grad_softmax", "n1"}, {{"T", "$T"}}},
  + {{"grad_x"}, "Mul", {"n2", "softmax"}, {{"T", "$T"}}}
  });
  // clang-format on
  return Status::OK();
  }
  REGISTER_OP_GRADIENT("Softmax", SoftmaxGrad);
  +Status LogSoftmaxGrad(const AttrSlice& attrs, FunctionDef* g) {
  + // clang-format off
  + *g = FDH::Define(
  + "LogSoftmaxGrad",
  + // Arg defs
  + {"x: T", "grad_logsoftmax: T"},
  + // Ret val defs
  + {"grad_x: T"},
  + // Attr defs
  + {{"T: {float, double}"}},
  + // Nodes
  + // Based on _LogSoftmaxGrad in nn_grad.py.
  + {
  + {{"softmax"}, "Softmax", {"x"}, {{"T", "$T"}}},
  + FDH::Const<int32>("indices", {-1}),
  + {{"n0"}, "Sum", {"grad_logsoftmax", "indices"},
  + {{"keep_dims", true}, {"T", "$T"}}},
  + {{"n1"}, "Mul", {"n0", "softmax"}, {{"T", "$T"}}},
  + {{"grad_x"}, "Sub", {"grad_logsoftmax", "n1"}, {{"T", "$T"}}}
  + });
  + // clang-format on
  + return Status::OK();
  +}
  +REGISTER_OP_GRADIENT("LogSoftmax", LogSoftmaxGrad);
  +
  Status ReluGrad(const AttrSlice& attrs, FunctionDef* g) {
  // clang-format off
  *g = FDH::Define(
  