1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_C_EAGER_MNIST_GRADIENTS_TESTUTIL_H_ 16 #define TENSORFLOW_C_EAGER_MNIST_GRADIENTS_TESTUTIL_H_ 17 #include <memory> 18 19 #include "absl/types/span.h" 20 #include "tensorflow/c/eager/abstract_tensor_handle.h" 21 #include "tensorflow/c/eager/c_api_experimental.h" 22 #include "tensorflow/c/eager/c_api_unified_experimental.h" 23 #include "tensorflow/c/eager/c_api_unified_experimental_internal.h" 24 #include "tensorflow/c/eager/gradients.h" 25 #include "tensorflow/c/eager/gradients_internal.h" 26 #include "tensorflow/c/experimental/ops/array_ops.h" 27 #include "tensorflow/c/experimental/ops/math_ops.h" 28 #include "tensorflow/c/experimental/ops/nn_ops.h" 29 #include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h" 30 #include "tensorflow/core/platform/status.h" 31 32 namespace tensorflow { 33 namespace gradients { 34 namespace internal { 35 36 // Computes 37 // y = inputs[0] + inputs[1] 38 // return grad(y, {inputs[0], inputs[1]}) 39 Status AddGradModel(AbstractContext* ctx, 40 absl::Span<AbstractTensorHandle* const> inputs, 41 absl::Span<AbstractTensorHandle*> outputs, 42 const GradientRegistry& registry); 43 44 // Computes 45 // y = inputs[0] * inputs[1] 46 // return grad(y, {inputs[0], inputs[1]}) 47 Status MatMulGradModel(AbstractContext* ctx, 48 absl::Span<AbstractTensorHandle* const> inputs, 49 absl::Span<AbstractTensorHandle*> outputs, 50 const GradientRegistry& registry); 51 52 // Computes 2-layer Neural Network with Softmax Loss. 53 Status MNISTForwardModel(AbstractContext* ctx, 54 absl::Span<AbstractTensorHandle* const> inputs, 55 absl::Span<AbstractTensorHandle*> outputs, 56 const GradientRegistry& registry); 57 58 // Computes MatMul with first matrix tranposed. 59 Status MatMulTransposeModel(AbstractContext* ctx, 60 absl::Span<AbstractTensorHandle* const> inputs, 61 absl::Span<AbstractTensorHandle*> outputs, 62 const GradientRegistry& registry); 63 64 // Test Model to verify Multi-grad functionality for MNIST 65 Status MNISTGradModel(AbstractContext* ctx, 66 absl::Span<AbstractTensorHandle* const> inputs, 67 absl::Span<AbstractTensorHandle*> outputs, 68 const GradientRegistry& registry); 69 70 // Test Model to verify scalar-tensor multiplication Op 71 Status ScalarMulModel(AbstractContext* ctx, 72 absl::Span<AbstractTensorHandle* const> inputs, 73 absl::Span<AbstractTensorHandle*> outputs, 74 const GradientRegistry& registry); 75 76 Status MatMulModel(AbstractContext* ctx, 77 absl::Span<AbstractTensorHandle* const> inputs, 78 absl::Span<AbstractTensorHandle*> outputs, 79 const GradientRegistry& registry); 80 81 Status MulModel(AbstractContext* ctx, 82 absl::Span<AbstractTensorHandle* const> inputs, 83 absl::Span<AbstractTensorHandle*> outputs, 84 const GradientRegistry& registry); 85 86 } // namespace internal 87 } // namespace gradients 88 } // namespace tensorflow 89 90 #endif // TENSORFLOW_C_EAGER_MNIST_GRADIENTS_TESTUTIL_H_ 91