TensorFlow.NET/src/TensorFlowNET.Core/Graphs/FuncGraph.cs

using Google.Protobuf;
using System;
using System.Collections.Generic;
using System.Linq;
using Tensorflow.Eager;
using Tensorflow.Exceptions;
using static Tensorflow.Binding;

namespace Tensorflow.Graphs
{
    /// <summary>
    /// Graph representing a function body.
    /// </summary>
    public class FuncGraph : Graph
    {
        IntPtr _func_graph_handle;
        public string FuncName => _graph_key;

        public Tensors Inputs { get; set; } = new Tensors();
        public Tensors Outputs { get; set; } = new Tensors();
        public Dictionary<string, string> Attrs { get; set; }

        Dictionary<long, (Tensor, Tensor)> _captures
            = new Dictionary<long, (Tensor, Tensor)>();

        public Tensor[] external_captures
            => _captures.Select(x => x.Value.Item1).ToArray();
        public (Tensor, Tensor)[] captures
            => _captures.Values.Select(x => x).ToArray();

        public Tensor[] internal_captures
            => _captures.Select(x => x.Value.Item2).ToArray();

        public Tensor[] captured_inputs
            => external_captures;

        /// <summary>
        /// Construct a new FuncGraph.
        /// </summary>
        public FuncGraph(string name) : base()
        {
            outer_graph = ops.get_default_graph();
            while (outer_graph.building_function)
                outer_graph = outer_graph.OuterGraph;
            _graph_key = name;
            building_function = true;
        }

        public FuncGraph(IntPtr handle, string name, Dictionary<string, string> attrs) : base()
        {
            outer_graph = ops.get_default_graph();
            while (outer_graph.building_function)
                outer_graph = outer_graph.OuterGraph;
            _graph_key = name;
            building_function = true;
            Attrs = attrs;
            // Will to test if FuncGraph has memory leak
            // c_api.TF_DeleteGraph(_handle);
            _handle = handle;
        }

        public void ToGraph(Operation[] opers,
            Tensor[] inputs, Tensor[] outputs,
            string[] output_names)
        {
            var status = new Status();
            _func_graph_handle = c_api.TF_GraphToFunction(_handle,
                _graph_key,
                false,
                opers.Length,
                opers.Select(x => (IntPtr)x).ToArray(),
                inputs.Length,
                inputs.Select(x => new TF_Output(x.op, 0)).ToArray(),
                outputs.Length,
                outputs.Select(x => new TF_Output(x.op, 0)).ToArray(),
                output_names == null || output_names.Length == 0 ? null : output_names,
                IntPtr.Zero,
                null,
                status.Handle);
            status.Check(true);

            SetAttrs();

            // c_api.TF_GraphCopyFunction(outer_graph, _func_graph_handle, IntPtr.Zero, status.Handle);
            // status.Check(true);

            c_api.TFE_ContextAddFunction(tf.Context.Handle, _func_graph_handle, status.Handle);
            status.Check(true);

            _graph_key = c_api.StringPiece(c_api.TF_FunctionName(_func_graph_handle));

            Inputs = inputs;
            // mark_as_return
            Outputs = outputs;// .Select(x => array_ops.identity(x)).ToArray();
        }

        public override Operation create_op(string op_type, Tensor[] inputs, TF_DataType[] dtypes, TF_DataType[] input_types = null, string name = null, Dictionary<string, AttrValue> attrs = null, OpDef op_def = null, bool compute_device = true)
        {
            foreach(var (i, inp) in enumerate(inputs))
                inputs[i] = capture(inp);

            return base.create_op(op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device);
        }

        const int _EAGER_CONST_THRESHOLD = 128;
        public Tensor capture(Tensor tensor, string name = null, Shape shape = null)
        {
            if(tensor is EagerTensor)
            {
                if (name == null)
                    name = ops.uid().ToString();

                // Small EagerTensors are captured with Const ops
                if (dtypes.is_value_dtype(tensor.dtype) 
                    && (tensor.rank == 0 || tensor.size < _EAGER_CONST_THRESHOLD))
                    return capture_eager_tensor(tensor, name);

                // Large EagerTensors and resources are captured with Placeholder ops
                return _capture_helper(tensor, name, shape: shape);
            }

            if(tensor.graph != this)
            {
                if (name == null)
                    name = tensor.op.name;
                var inner_graph = tensor.graph;
                while(inner_graph != null && inner_graph is FuncGraph inner_func_graph)
                {
                    if (inner_graph == this)
                        throw new InaccessibleTensorError($"The tensor '{tensor.name}' cannot be accessed here: it is defined" +
                            " in another function or code block. Use return values," +
                            " explicit Python locals or TensorFlow collections to access" +
                            $" it. Defined in: {tensor.graph.graph_key}; accessed from: {graph_key}.");
                    inner_graph = inner_func_graph.outer_graph;
                }
                return _capture_helper(tensor, name);
            }

            return tensor;
        }

        Tensor capture_eager_tensor(Tensor tensor, string name)
        {
            Tensor graph_const = null;
            if (!_captures.ContainsKey(tensor.Id))
            {
                graph_const = tf_with(ops.control_dependencies(null), ctl
                    => constant_op.constant(tensor.numpy(), dtype: tensor.dtype, shape: tensor.shape, name: name));
                add_capture(tensor, graph_const);
            }
            else
            {
                graph_const = _captures[tensor.Id].Item2;
            }

            BackwardFunction _backward_function_wrapper = (output_grads, unneeded_gradients) =>
            {
                return output_grads;
            };

            tf.Runner.RecordGradient("captured_value",
                new[] { graph_const }, null,
                new[] { tensor },
                getBackwardFunction: () => _backward_function_wrapper
                /*getForwardFunction: forward_function*/);

            return graph_const;
        }

        Tensor _capture_helper(Tensor tensor, string name, Shape shape = null)
        {
            Tensor placeholder = null;
            if (!_captures.ContainsKey(tensor.Id))
            {
                placeholder = _create_substitute_placeholder(tensor,
                    name: name,
                    dtype: tensor.dtype,
                    shape: shape);
                add_capture(tensor, placeholder);
            }
            else
            {
                placeholder = _captures[tensor.Id].Item2;
            }

            BackwardFunction _backward_function_wrapper = (output_grads, unneeded_gradients) =>
            {
                return output_grads;
            };

            tf.Runner.RecordGradient("captured_value",
                new[] { placeholder }, null,
                new[] { tensor },
                getBackwardFunction: () => _backward_function_wrapper
                /*getForwardFunction: forward_function*/);

            return placeholder;
        }

        void add_capture(Tensor tensor, Tensor placeholder)
        {
            _captures.Add(tensor.Id, (tensor, placeholder));
            Inputs.Add(placeholder);
        }

        Tensor _create_substitute_placeholder(Tensor value, 
            string name = null, 
            TF_DataType dtype = TF_DataType.DtInvalid, 
            Shape shape = null)
        {
            if (shape is null)
                shape = value.shape;
            if (dtype == TF_DataType.DtInvalid)
                dtype = value.dtype;

            var placeholder = tf_with(ops.control_dependencies(null), ctl
                => array_ops.placeholder(dtype, shape: shape, name: name));
            // custom_gradient.copy_handle_data(value, placeholder)
            return placeholder;
        }

        void SetAttrs()
        {
            if (Attrs == null)
                return;

            foreach (var (_name, attr_value) in enumerate(Attrs))
            {
                var serialized = new AttrValue
                {
                    S = ByteString.CopyFromUtf8(attr_value)
                }.ToByteArray();
                c_api.TF_FunctionSetAttrValueProto(_func_graph_handle, _name, serialized, serialized.Length, tf.Status.Handle);
                tf.Status.Check(true);
            }
        }

        public override Graph as_default()
        {
            tf.Context.graph_mode(isFunc: true);
            ops.set_default_graph(this);
            return this;
        }

        public override void Exit()
        {
            tf.Context.restore_mode();
            ops.pop_graph();
        }

        protected override void DisposeUnmanagedResources(IntPtr handle)
        {
            c_api.TFE_ContextRemoveFunction(tf.Context.Handle, _graph_key, tf.Status.Handle);
            c_api.TF_DeleteFunction(_func_graph_handle);
            base.DisposeUnmanagedResources(handle);
        }
    }
}