3c54ef5ab9
- Support packing multiple EagerTensors of the same dtype and shape. - Create CompositeDevices on the same task as the local host CPU, in order to correctly trigger packed TensorHandle copy from a client to a remote worker. PiperOrigin-RevId: 312164194 Change-Id: Ia15718309c8c68eb645bfe0bf967ddd6d2551b3a
398 lines
18 KiB
C
Executable File
398 lines
18 KiB
C
Executable File
/* 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/lib/core/status.h"
|
|
#include "tensorflow/core/lib/gtl/inlined_vector.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 TFE Context
|
|
// Item 2: device_name: Name of the device on which to execute the operation,
|
|
// or NULL for automatic selection.
|
|
// Item 3: op_name: Name of the TensorFlow op to execute.
|
|
// Item 4: name: An optional name for the operation.
|
|
// Item 5: List representing all callbacks to execute after successful
|
|
// op execute.
|
|
// Item 6 onwards: inputs - This is a list of inputs followed by a list of
|
|
// attrs. It is not necessary for type attrs to be present.
|
|
//
|
|
// 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();
|
|
|
|
// 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);
|
|
#endif // TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_
|