Tensorflow Integration

Save model

import tensorflow as tf

class Adder(tf.Module):
@tf.function
def add(self, x):
    return x + x

model = Adder()

tf.saved_model.save(model, "./", signatures=model.add.get_concrete_function(tf.TensorSpec([], tf.float32)))

Use matxscript load SavedModel

import matx

tf_op = matx.script("./", backend='TensorFlow', device=-1, use_xla=0, allow_growth=False)

Trace and inference

ix = matx.NDArray([1], [1], 'float32')

def process(x):
    return tf_op({"x":x})

ret = process(ix)
print(ret)

s = matx.trace(process, ix)
ret = s.run({"x":ix})
print(ret)