CODEWITHSUNDEEP
pythontensorflow
p.koch
p.koch

Reputation: 183

What Python types does TensorFlow accept for Attr's of type "tensor"?

I am defining a new Op in C++ which takes in a single attribute of type tensor, roughly following these instructions. A stripped version of the Op code is below:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

REGISTER_OP("DoStuff")
    .Attr("attr: tensor = { dtype: DT_FLOAT }")
    .Input("in: float")
    .Output("out: float");

class DoStuffOp : public OpKernel {
public:
    explicit DoStuffOp(OpKernelConstruction *context) : OpKernel(context) {
        OP_REQUIRES_OK(context, context->GetAttr("attr", &attr_));
        // ...
    }

    void Compute(OpKernelContext *context) override {
        // ...
    }

private:
    Tensor attr_;
};

REGISTER_KERNEL_BUILDER(Name("DoStuff").Device(DEVICE_CPU), DoStuffOp);

I can compile the Op into a .so file fine. However, I can't figure out how to successfully pass in a value for attr. When I run the following in Python:

import tensorflow as tf
dostufflib = tf.load_op_library('build/do_stuff.so')
sess = tf.InteractiveSession()

A = [[1.0, 2.0, 3.0],
     [4.0, 5.0, 6.0]]
X = tf.Variable(tf.constant(1.0))

Y = dostufflib.do_stuff(X, A)

I get TypeError: Don't know how to convert [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] to a TensorProto for argument 'attr'. Nothing I do seems to satisfy the type conversion: list, numpy array, tf.Tensor, tf.Variable, etc. How do you pass Python variables into an Op as tensor attributes?

Upvotes: 1

Views: 393

Answers (1)

p.koch
p.koch

Reputation: 183

After much more hunting, I found tf.contrib.util.make_tensor_proto, a function that converts a python scalar, python list, numpy ndarray, or numpy scalar into a tf.TensorProto object. The following works:

A = tf.contrib.util.make_tensor_proto([[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]])
X = tf.Variable(tf.constant(1.0))

Y = dostufflib.do_stuff(X, A)

Upvotes: 3

Related Questions