init

drop changes

drop changes

drop changes

drop changes

drop changes

Author: MarisaKirisame

python/tvm/relay/op/_tensor_grad.py

@@ -320,6 +320,18 @@ def dense_grad(orig, grad):
     return [collapse_sum_like(transpose(grad) * weight, data),
             collapse_sum_like(data * transpose(grad), weight)]
 
+# UNTESTED
+@register_gradient("take")
+def take_grad(orig, grad):
+    x, y = orig.args
+    return [zeros_like(x), zeros_like(y)]
+    return [Call(op_get("take_grad"), [x, y, grad], orig.attrs), zeros_like(y)]
+
+
+@register_gradient("shape_of")
+def shape_of_grad(orig, grad):
+    return [zeros_like(orig.args[0])]
+
 
 @register_gradient("reshape")
 def reshape_grad(orig, grad):
@@ -365,12 +377,3 @@ def sum_grad(orig, grad):
     """Returns grad broadcasted to data dims"""
     data = orig.args[0]
     return [broadcast_to_like(grad, data)]
-
-
-@register_gradient("nn.cross_entropy")
-def cross_entropy_grad(orig, grad):
-    x, y = orig.args
-    shape = shape_of(x)
-    batch_size = take(shape, const(0, dtype='int32'), axis=0)
-    grad = grad / batch_size.astype('float32')
-    return [-grad * y / x, -grad * log(x)]

python/tvm/relay/op/transform.py

@@ -42,12 +42,15 @@ def cast(data, dtype):
 
 def cast_like(data, dtype_like):
     """Cast input tensor to data type of another tensor.
+
     Parameters
     ----------
     data : relay.Expr
         The input data to the operator.
+
     dtype_like: relay.Expr
         The tensor to cast to.
+
     Returns
     -------
     result : relay.Expr
@@ -249,8 +252,8 @@ def reshape_like(data, shape_like):
     data : relay.Expr
         The input data to the operator.
 
-    shape_like : tuple of int
-        The new shape. Should be compatible with the original shape.
+    shape_like : relay.Expr
+        The tensor to reshape to. Should be compatible with the original shape.
 
     Returns
     -------
@@ -260,6 +263,37 @@ def reshape_like(data, shape_like):
     return _make.reshape_like(data, shape_like)
 
 
+def embed_like(data, indices, type_like, axis=None, mode="clip"):
+    """Take elements from an array along an axis.
+
+    Parameters
+    ----------
+    data : relay.Expr
+        The source array.
+
+    indices : rely.Expr
+        The indices of the values to extract.
+
+    type_like : relay.Expr
+        The tensor that provide the type to embed into.
+
+    axis : int, optional
+        The axis over which to select values. By default,
+        the flattened input array is used.
+
+    mode : str, optional
+        Specifies how out-of-bound indices will behave [clip, wrap, fast].
+        clip: clip to the range (default).
+        wrap: wrap around the indices.
+        fast: no clip or wrap around (user must make sure indices are in-bound).
+
+    Returns
+    -------
+    ret : relay.Expr
+        The computed result.
+    """
+    return _make.embed_like(data, indices, type_like, axis, mode)
+
 def take(data, indices, axis=None, mode="clip"):
     """Take elements from an array along an axis.
 

src/relay/op/tensor/transform.cc

@@ -187,6 +187,7 @@ RELAY_REGISTER_OP("reinterpret")
     .set_attr<TOpPattern>("TOpPattern", kElemWise)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout);
 
+
 // relay.expand_dims
 TVM_REGISTER_NODE_TYPE(ExpandDimsAttrs);
 
@@ -958,6 +959,43 @@ Examples::
 .set_attr<FTVMCompute>("FTVMCompute", TakeCompute)
 .set_attr<TOpPattern>("TOpPattern", kInjective);
 
+bool EmbedLikeRel(const Array<Type>& types,
+                  int num_inputs,
+                  const Attrs& attrs,
+                  const TypeReporter& reporter) {
+  // `types` contains: [data, indices, type_like, result]
+  CHECK_EQ(types.size(), 4);
+  reporter->Assign(types[3], types[2]);
+  return TakeRel({types[2], types[1], types[0]}, 2, attrs, reporter);
+}
+
+Expr MakeEmbedLike(Expr data,
+                   Expr indices,
+                   Expr type_like,
+                   Integer axis,
+                   std::string mode) {
+  auto attrs = make_node<TakeAttrs>();
+  attrs->axis = std::move(axis);
+  attrs->mode = std::move(mode);
+  static const Op& op = Op::Get("embed_like");
+  return CallNode::make(op, {data, indices, type_like}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op._make.embed_like")
+.set_body_typed(MakeEmbedLike);
+
+RELAY_REGISTER_OP("embed_like")
+.describe(R"code(The inverse of take.)code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.TakeAttrs")
+.set_num_inputs(3)
+.add_argument("data", "Tensor", "The input tensor.")
+.add_argument("indices", "Tensor", "The indices tensor.")
+.add_argument("type_like", "Tensor", "The tensor that provide the type and shape to embed into.")
+.set_support_level(3)
+.add_type_rel("EmbedLike", EmbedLikeRel)
+.set_attr<FTVMCompute>("FTVMCompute", TakeCompute) // implement this at python side?
+.set_attr<TOpPattern>("TOpPattern", kInjective);
+
 
 // Init ops
 TVM_REGISTER_NODE_TYPE(InitOpAttrs);

tests/python/relay/test_pass_gradient.py

@@ -23,9 +23,14 @@
 from tvm.relay.transform import gradient
 from tvm.relay.prelude import Prelude
 from tvm.relay.testing import add_nat_definitions, make_nat_expr, run_infer_type, check_grad, rand
+from tvm.relay.testing import resnet, inception_v3, squeezenet, densenet, lstm
 import tvm.relay.op as op
 
 
+def rand(dtype='float32', *shape):
+    return tvm.nd.array(np.random.rand(*shape).astype(dtype))
+
+
 def test_id():
     shape = (10, 10)
     dtype = 'float32'
@@ -198,10 +203,9 @@ def test_pow():
     double = relay.Function([x], x + x)
     i = relay.var("i", t)
     func = relay.Function([i], p.nat_iterate(double, make_nat_expr(p, 3))(i))
-    mod["main"] = func
-    mod["main"] = gradient(mod["main"], mod=mod)
-    m = transform.InferType()(mod)
-    back_func = m["main"]
+    mod["func"] = func
+    mod["back_func"] = gradient(mod["func"], mod=mod)
+    back_func = mod["back_func"]
     assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])]))
     i_nd = rand(dtype, *shape)
     ex = create_executor(mod=mod)
@@ -295,6 +299,19 @@ def test_concat():
     # no value validation as concatenate has dummy gradient right now.
 
 
+def rand_from_type(t):
+    assert isinstance(t, relay.ty.TensorType)
+    return rand(t.dtype, *[int(s) for s in t.shape])
+
+
+def test_resnet():
+    x, _ = densenet.get_workload()
+    x = gradient(x["main"])
+    args = [rand_from_type(e.checked_type) for e in x.params]
+    ex = create_executor()
+    ex.evaluate(x)(*args)
+
+
 if __name__ == "__main__":
     test_id()
     test_add()