Repository: apache/tvm

Issue:
[DSL/TE] Scalar support for `te.extern`
We have scalar support in `te.compute` now. But for `te.extern`, tvm will raise error for scalar input. It would be great that we can have consistent behavior.

Sample code:
```
import tvm
from tvm import te
import numpy as np

a = np.array(1.0, 'float32')
b = np.array(0.0, 'float32')

@tvm.register_func("tvm.copy")
def mycopy(x, y):
    x.copyto(y)

A = te.placeholder(a.shape, name='A')
B = te.extern(a.shape, [A],
              lambda ins, outs: tvm.tir.call_packed(
              "tvm.copy", ins[0], outs[0]), name="B")

s = te.create_schedule(B.op)
f = tvm.build(s, [A, B], 'llvm')

ta = tvm.nd.array(a)
tb = tvm.nd.array(b)
f(ta, tb)
print(tb)
```


Discussion:

Resolving patch:
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -208,7 +208,7 @@ def extern(shape,
            out_buffers=None,
            tag="",
            attrs=None):
-    """Compute several tensor via extern function.
+    """Compute several tensors via an extern function.
 
     Parameters
     ----------
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -28,7 +28,6 @@
 def _pack_buffer(buf):
     """Build intrinsics that packs the buffer.
     """
-    assert buf.shape
     shape = Call("handle", "tir.tvm_stack_make_shape", buf.shape)
     strides = Call("handle", "tir.tvm_stack_make_shape", buf.strides) if buf.strides else 0
     pack_args = [buf.data,


Test patch:
diff --git a/tests/python/unittest/test_te_tensor.py b/tests/python/unittest/test_te_tensor.py
--- a/tests/python/unittest/test_te_tensor.py
+++ b/tests/python/unittest/test_te_tensor.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import tvm
+import numpy as np
 from tvm import te
 from topi.nn.pooling import pool
 
@@ -303,6 +304,52 @@ def intrin_func(ins, outs):
     s[P].tensorize(oh, intrin)
     tvm.lower(s, [A, P])
 
+def test_tensor_scalar_mixed():
+    # test te with tensor and scalar
+    a = np.array(np.random.uniform(size=(10,)), 'float32')
+    b = np.array(np.random.uniform(size=(1))[0], 'float32')
+    c = np.array(np.random.uniform(size=(10,)), 'float32')
+
+    @tvm.register_func("tvm.test_tensor_scalar_scale")
+    def my_scale(tensor, scalar, out):
+        out_np = tensor.asnumpy() * scalar.asnumpy()
+        tvm.nd.array(out_np).copyto(out)
+
+    A = te.placeholder(a.shape, name='A')
+    B = te.placeholder(b.shape, name='B')
+    C = te.extern(a.shape, [A, B],
+                  lambda ins, outs: tvm.tir.call_packed(
+                  "tvm.test_tensor_scalar_scale", ins[0], ins[1], outs[0]), name="C")
+    s = te.create_schedule(C.op)
+    f = tvm.build(s, [A, B, C], 'llvm')
+
+    ta = tvm.nd.array(a)
+    tb = tvm.nd.array(b)
+    tc = tvm.nd.array(c)
+    f(ta, tb, tc)
+    tvm.testing.assert_allclose(a * b, tc.asnumpy())
+
+
+def test_tensor_scalar():
+    # test te with scalar shape
+    a = np.array(np.random.uniform(size=(1))[0], 'float32')
+    b = np.array(0.0, 'float32')
+
+    @tvm.register_func("tvm.test_tensor_scalar_copy")
+    def mycopy(x, y):
+        x.copyto(y)
+
+    A = te.placeholder(a.shape, name='A')
+    B = te.extern(a.shape, [A],
+                  lambda ins, outs: tvm.tir.call_packed(
+                  "tvm.test_tensor_scalar_copy", ins[0], outs[0]), name="B")
+    s = te.create_schedule(B.op)
+    f = tvm.build(s, [A, B], 'llvm')
+
+    ta = tvm.nd.array(a)
+    tb = tvm.nd.array(b)
+    f(ta, tb)
+    tvm.testing.assert_allclose(ta.asnumpy(), tb.asnumpy())
 
 if __name__ == "__main__":
     test_rank_zero()
@@ -321,3 +368,5 @@ def intrin_func(ins, outs):
     test_tuple_inputs()
     test_tuple_with_different_deps()
     test_tensor_pool()
+    test_tensor_scalar()
+    test_tensor_scalar_mixed()