xref: /external/tensorflow/tensorflow/python/eager/pywrap_tfe.h

/* Copyright 2017 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.
==============================================================================*/

#ifndef TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_
#define TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_

// Place `<locale>` before <Python.h> to avoid build failure in macOS.
#include <locale>

// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>

#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"

typedef tensorflow::gtl::InlinedVector<TFE_TensorHandle*, 4>
    TFE_InputTensorHandles;
typedef tensorflow::gtl::InlinedVector<TFE_TensorHandle*, 2>
    TFE_OutputTensorHandles;

// Execute a TensorFlow operation.
//
// 'device_name': Name of the device on which to execute the operation, or NULL
//                for automatic selection.
// 'op_name': Name of the TensorFlow op to execute.
// 'inputs': An array of TFE_TensorHandle*'s of size 'num_inputs'. These tensors
//           will be provided as input to the operation.
// 'attrs': A Python tuple alternating names and attr values.
// 'outputs': A pointer to a TFE_OutputTensorHandles in which outputs will
//            placed. On success, its elements will be filled in and the
//            caller takes ownership of each returned TFE_TensorHandle.
//            'outputs' MUST be sized to be at least as large as the number
//            of tensors produced by the operation and will be resized to
//            the actual number of tensors produced.
void TFE_Py_Execute(TFE_Context* ctx, const char* device_name,
                    const char* op_name, TFE_InputTensorHandles* inputs,
                    PyObject* attrs, TFE_OutputTensorHandles* outputs,
                    TF_Status* out_status);

// Execute a cancelable TensorFlow operation.
//
// Arguments as above (for TFE_Py_Execute), with the addition of:
// 'cancellation_manager': A pointer to a TFE_CancellationManager that can be
//                         used to cancel execution of the given operation.
typedef struct TFE_CancellationManager TFE_CancellationManager;
void TFE_Py_ExecuteCancelable(TFE_Context* ctx, const char* device_name,
                              const char* op_name,
                              TFE_InputTensorHandles* inputs, PyObject* attrs,
                              TFE_CancellationManager* cancellation_manager,
                              TFE_OutputTensorHandles* outputs,
                              TF_Status* out_status);

// Registers e as the Exception class for handling not ok Status. Returns
// Py_None if registration succeeds, else throws a TypeError and returns NULL.
//
// This function is not thread-safe.
PyObject* TFE_Py_RegisterExceptionClass(PyObject* e);

// Registers e as the VSpace to use.
// `vspace` must be a imperative_grad.py:VSpace named tuple.
PyObject* TFE_Py_RegisterVSpace(PyObject* e);

// Registers e as the Exception to be raised when the conditions of
// TFE_Py_FastPathExecute_C have not been met. When this exception is set, it
// is a signal to the calling code that it should fall back to the safer (and
// more complete) code path.
//
// This function is not thread-safe.
PyObject* TFE_Py_RegisterFallbackExceptionClass(PyObject* e);

// Registers e as the gradient_function.
// The registered function takes
// (op_name, attrs, num_inputs, inputs, outputs, output_gradients) and returns
// the input gradients. This function will not correctly be able to generate
// gradients for functional ops - the gradients for those ops are calculated
// through a different codepath (see function.py for additional information).
//
// This function is not thread-safe.
PyObject* TFE_Py_RegisterGradientFunction(PyObject* e);

// Registers e as the forward_gradient_function.  The registered function takes
// (op_name, attrs, inputs, outputs, tangents) and returns the output
// tangents. This function is used only for operations, not for custom gradients
// or functional ops.
//
// This function is not thread-safe.
PyObject* TFE_Py_RegisterJVPFunction(PyObject* e);

// Returns 0 if 'status' is TF_OK. Otherwise, raises an exception (using
// `exception` if not nullptr, else using the class registered via
// TFE_Py_RegisterExceptionClass), and returns -1.
int MaybeRaiseExceptionFromTFStatus(TF_Status* status, PyObject* exception);

// Returns 0 if 'status' is ok. Otherwise, raises an exception (using
// `exception` if not nullptr, else using the class registered via
// TFE_Py_RegisterExceptionClass), and returns -1.
int MaybeRaiseExceptionFromStatus(const tensorflow::Status& status,
                                  PyObject* exception);

// Returns the string associated with the passed-in python object.
const char* TFE_GetPythonString(PyObject* o);

// Returns a unique id on each call.
int64_t get_uid();

// Wraps the output of get_uid as a Python Long object. Ownership is passed to
// the caller.
PyObject* TFE_Py_UID();

// Deleter for Context objects, called from the Capsule that owns it.
void TFE_DeleteContextCapsule(PyObject* context);

// Returns true if o is an instance of EagerTensor, but not a subclass. Else
// returns false.
bool EagerTensor_CheckExact(const PyObject* o);

// Helper function to construct a new EagerTensor from a TFE_TensorHandle.
PyObject* EagerTensorFromHandle(TFE_TensorHandle* handle,
                                const bool is_packed = false);

// Extracts the handle inside EagerTensor object `o`. Returns nullptr on error.
TFE_TensorHandle* EagerTensor_Handle(const PyObject* o);

// Creates the `EagerTensor` class by subclassing `base_class` and returns the
// newly created type, or nullptr on error.
PyObject* TFE_Py_InitEagerTensor(PyObject* base_class);

// Sets `profiler` as the current profiler to receive callbacks about events
// on eager tensors. Currently, the only reported event is creation.
// `profiler` is expected to have a `created(self, eager_tensor)` method that
// takes the created tensor as its single argument.
// Previous profiler, if any, is unset and will not receive any more
// callbacks.
// To unset the profiler, pass Py_None as the value of `profiler`.
PyObject* TFE_Py_SetEagerTensorProfiler(PyObject* profiler);

// Creates a new tape and adds it to the active set. `persistent` and
// `watch_accessed_variables` must be `PyBool_Type` (`Py_True` or `Py_False`).
PyObject* TFE_Py_TapeSetNew(PyObject* persistent,
                            PyObject* watch_accessed_variables);

// Removes the passed tape from the set of active tapes.
void TFE_Py_TapeSetRemove(PyObject* tape);

// Adds the passed tape to the set of active tapes.
void TFE_Py_TapeSetAdd(PyObject* tape);

// Returns true if the tape stack is empty.
PyObject* TFE_Py_TapeSetIsEmpty();

// Check if any backward tape should record an operation given inputs.
//
// Does not take forward accumulators into account.
PyObject* TFE_Py_TapeSetShouldRecordBackprop(PyObject* tensors);

// Determine possible gradient types, taking forward accumulators into account.
//   - 0 if no tape will record (implies TFE_Py_TapeSetShouldRecordBackprop
//     is false and no forward accumulator is watching)
//   - 1 if first-order gradients may be requested
//   - 2 if higher-order gradients may be requested
PyObject* TFE_Py_TapeSetPossibleGradientTypes(PyObject* tensors);

void TFE_Py_TapeWatch(PyObject* tape, PyObject* tensor);
void TFE_Py_TapeSetDeleteTrace(tensorflow::int64 tensor_id);

// Stops any gradient recording on the current thread.
//
// Includes forward accumulators.
void TFE_Py_TapeSetStopOnThread();

// Restarts gradient recording on the current thread.
void TFE_Py_TapeSetRestartOnThread();

// Checks whether gradient recording is stopped on the current thread.
PyObject* TFE_Py_TapeSetIsStopped();

// Records an operation for the purpose of gradient computation.
//
// Arguments:
//  - op_type is a string for the operation type, used in the backprop code
//  - output_tensors are a list of Python Tensor objects output by the operation
//  - input_tensors are a list of input Tensors to the recorded operation
//  - backward_function is the function to be called during backprop or
//    forwardprop to, given the gradients of the output tensors, produce the
//    gradients of the input tensors. This function is automatically transposed
//    during forwardprop.
//  - forward_function is an optional special-case for forwardprop, taking input
//    jvps and returning output jvps.
//
// Records an operation both for backprop (gradient tape) and forwardprop
// (forward accumulator). Equivalent to calling both
// TFE_Py_TapeSetRecordOperationBackprop and
// TFE_Py_TapeSetRecordOperationForwardprop.
PyObject* TFE_Py_TapeSetRecordOperation(PyObject* op_type,
                                        PyObject* output_tensors,
                                        PyObject* input_tensors,
                                        PyObject* backward_function,
                                        PyObject* forward_function);

// Records an operation only for backprop (gradient tapes).
//
// Same arguments as TFE_Py_TapeSetRecordOperation.
PyObject* TFE_Py_TapeSetRecordOperationBackprop(PyObject* op_type,
                                                PyObject* output_tensors,
                                                PyObject* input_tensors,
                                                PyObject* backward_function);

// Records an operation only for forwardprop (forward accumulators).
//
// Arguments:
//  - op_type is a string for the operation type, used in the backprop code
//  - output_tensors are a list of Python Tensor objects output by the operation
//  - input_tensors are a list of input Tensors to the recorded operation
//  - backward_function is the function to be called to, given the gradients of
//    the output tensors, produce the gradients of the input tensors. This
//    function is automatically transposed to produce output gradients given
//    input gradients.
//  - forwardprop_output_indices indicates any output_tensors which contain
//    JVPs. Typically these will have come from TFE_Py_PackJVPs. May
//    be None or an empty sequence if there are no JVP outputs from the
//    operation.
PyObject* TFE_Py_TapeSetRecordOperationForwardprop(
    PyObject* op_type, PyObject* output_tensors, PyObject* input_tensors,
    PyObject* backward_function, PyObject* forwardprop_output_indices);

// Notifies all tapes that a variable has been accessed.
void TFE_Py_TapeVariableAccessed(PyObject* variable);

// Watches the given variable object on the given tape.
void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable);

// Computes a gradient based on information recorded on the tape.`tape` must
// have been produced by TFE_Py_NewTape. `target` and `sources` must be python
// lists of Tensor objects. `output_gradients` is either None or a python list
// of either Tensor or None, and if not None should have the same length as
// target.
PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* target,
                              PyObject* sources, PyObject* output_gradients,
                              PyObject* sources_raw,
                              PyObject* unconnected_gradients,
                              TF_Status* status);

// Execute a tensorflow operation assuming that all provided inputs are
// correctly formatted (i.e. EagerTensors). If it doesn't find EagerTensors,
// it will simply fail with a NotImplementedError.
//
// The "args" PyObject* is meant to be a tuple with the following structure:
//  Item 1: The Python eager Context object
//  Item 2: op_name: Name of the TensorFlow op to execute.
//  Item 3: name: An optional name for the operation.
//  Item 4 onwards: inputs - This is a list of inputs followed by a list of
//        attrs. It is not necessary for type attrs to be present.
//
// Note: the device_name and op_callbacks, which were previously passed
// as arguments, are now read via GetEagerContextThreadLocalData().
//
// This is named _C since there doesn't seem to be any way to make it visible
// in the SWIG interface without renaming due to the use of the %native
// directive.
PyObject* TFE_Py_FastPathExecute_C(PyObject* args);

// Record the gradient for a given op.
PyObject* TFE_Py_RecordGradient(PyObject* op_name, PyObject* inputs,
                                PyObject* attrs, PyObject* results,
                                PyObject* forward_pass_name_scope);

// Returns all variables watched by the given tape in the order those variables
// were created.
PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape);

// Creates a new forward accumulator. Does not add it to the active set.
PyObject* TFE_Py_ForwardAccumulatorNew(bool use_batch);

// Adds a ForwardAccumulator to the active set, meaning it will watch executed
// operations. It must not already be in the active set.
PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator);
// Removes a forward accumulator from the active set, meaning it will no longer
// be watching operations.
void TFE_Py_ForwardAccumulatorSetRemove(PyObject* accumulator);

// Tell the forward accumulator `accumulator` to watch `tensor`, with a Tensor
// tangent vector `tangent` of matching shape and dtype.
void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor,
                                    PyObject* tangent);

// Looks up the Jacobian-vector product of `tensor` in the forward accumulator
// `accumulator`. Returns None if no JVP is available.
PyObject* TFE_Py_ForwardAccumulatorJVP(PyObject* accumulator, PyObject* tensor);

// Temporarily push or pop transient state for accumulators in the active set.
//
// Allows an accumulator which is currently processing an operation to
// temporarily reset its state. This is useful when building forwardprop
// versions of functions, where an accumulator will trigger function building
// and then must process captured symbolic tensors while building it. Without
// pushing and popping, accumulators ignore operations executed as a direct
// result of their own jvp computations.
PyObject* TFE_Py_ForwardAccumulatorPushState();
PyObject* TFE_Py_ForwardAccumulatorPopState();

// Collects state from all current forward accumulators related to `tensors`.
//
// This is useful for packing JVPs as function inputs before executing a
// function which computes primals and JVPs at the same time.
//
// Does not include accumulators which are currently in the process of computing
// a jvp (and so appear somewhere on the current execution stack) or any
// accumulators more deeply nested.
//
// Includes JVPs for `tensors` and any higher-order JVPs for those
// (recursively). Returns a two-element tuple (indices, jvps):
//   indices: A sequence of sequences of two-element tuples. Each forward
//       accumulator is represented as a sequence of tuples with (primal_index,
//       jvp_index). Both integers index into the concatenated `tensors + jvps`
//       array.
//   jvps: A flat list of Tensors. Best interpreted as a sequence to be
//       appended to `tensors`.
PyObject* TFE_Py_PackJVPs(PyObject* tensors);

// Variable Watcher methods.

// Creates a new variable watcher and adds it to the set of active variable
// watchers.
PyObject* TFE_Py_VariableWatcherNew();

// Removes the passed variable watcher from the set of active variable watchers.
void TFE_Py_VariableWatcherRemove(PyObject* variable_watcher);

// Notifies all variable watchers that a variable has been accessed.
void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable);

// Returns all variables watched by the given variable_watcher in the order
// those variables were created.
PyObject* TFE_Py_VariableWatcherWatchedVariables(PyObject* variable_watcher);

// Returns an EagerTensor of dimension [len(`tensors`)] containing
// the `slice_dim`'th dimension of each tensor in `tensors`. In other words,
// TFE_Py_TensorShapeSlice takes a slice of dimensions of tensors in
// `tensors`. For example, if `tensors` contains tensors of with shapes
// [1, 2, 3], [4, 5], [6, 7, 8, 9], TFE_Py_TensorShapeSlice called with
// `slice_dim` equal to 1 will return [2, 5, 7].
// On error, returns nullptr and sets python exception.
// REQUIRES: `tensors` is a python list/tuple of EagerTensors
// REQUIRES: `slice_dim` is non-negative and smaller than the rank of all
//   tensors in `tensors`.
PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim);

// Returns the shape of this tensor's on-device representation.
// The shape is represented as a Python tuple of integers.
PyObject* TFE_Py_TensorShapeOnDevice(PyObject* tensor);

// Encodes the object as a tuple that is meant to be used as part of the key
// for the defun function cache.  If `include_tensor_ranks_only` is true,
// then the encoding only stores tensor ranks, and the key is
// agnostic to dimension sizes.  Otherwise, full tensor shape encodings are
// returned.
PyObject* TFE_Py_EncodeArg(PyObject*, bool include_tensor_ranks_only);

void TFE_Py_EnableInteractivePythonLogging();

// Sets `python_context` as the current eager Context object (defined
// in eager/context.py). This function must be called at least once before
// eager tensors are created.
// If an error is encountered, sets python error and returns NULL. Else, returns
// Py_None.
//
// This function is not thread-safe.
PyObject* TFE_Py_SetEagerContext(PyObject* py_context);

// Returns the current eager Context object (defined in eager/context.py)
// that was last set using TFE_Py_SetEagerContext.
// If an error is encountered, sets python error and returns NULL.
// The returned PyObject is "new", i.e. the caller must call Py_DECREF on it at
// some point.
PyObject* GetPyEagerContext();

// These are exposed since there is SWIG code that calls these.
// Returns a pre-allocated status if it exists.
TF_Status* GetStatus();
// Returns the pre-allocated status to the code.
void ReturnStatus(TF_Status* status);

namespace tensorflow {

// Returns the DataType for the specified tensor.  Returns DT_INVALID if
// PyObject is not a tensor.
DataType PyTensor_DataType(PyObject* tensor);

// Thread-local data associated with a Python eager Context object.
//
// TODO(edloper): Consider changing device_name and scope_name to a const char*
// (with nullptr used for None). However, note that existing code (e.g.
// TFE_TensorHandleCache::Lookup) assumes that the lifetime of these strings
// extends beyond the point where their value is changed; so we'd need to make
// sure that the strings stay alive (maybe using PyUnicode_InternInPlace?)
struct EagerContextThreadLocalData {
  bool is_eager = false;
  bool invoking_op_callbacks = false;
  tensorflow::Safe_PyObjectPtr device_name;
  tensorflow::Safe_PyObjectPtr scope_name;
  tensorflow::Safe_PyObjectPtr device_spec;
  tensorflow::Safe_PyObjectPtr function_call_options;
  tensorflow::Safe_PyObjectPtr executor;
  tensorflow::Safe_PyObjectPtr op_callbacks;
};

// Create a thread-local-data structure associated with py_eager_context.
// `is_eager` and `device_spec` are used to supply default values for those
// fields whenever a new thread-local instance is created for py_eager_tensor.
//
// This function assumes that the Python GIL is held (and does not perform its
// own locking).
void MakeEagerContextThreadLocalData(PyObject* py_eager_context,
                                     PyObject* is_eager,
                                     PyObject* device_spec);

// Returns the thread-local instance of EagerContextThreadLocalData that is
// associated with the given Python Context object.  If an instance has not
// yet been created for `py_eager_context` in this thread, then a new one is
// created, and initialized with the default values specified in
// MakeEagerContextThreadLocalData.
EagerContextThreadLocalData* GetEagerContextThreadLocalData(
    PyObject* py_eager_context);

// Free data structures used to track py_eager_context.
//
// This frees global state associated with py_eager_context, as well as thread-
// local state associated with py_eager_context and the current thread. If you
// wish to destroy thread-local state associated with a single py_eager_context
// for multiple threads, then you must call this method from each thread.
//
// Thread-local state assocaited with eager contexts is also automatically
// cleaned up when the thread is destroyed.
//
// This function assumes that the Python GIL is held (and does not perform its
// own locking).
void DestroyEagerContextThreadLocalData(PyObject* py_eager_context);

}  // namespace tensorflow

#endif  // TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_