(tvm) branch main updated: [BugFix]Ensure that bf16 arrays are created as expected (#16436)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new ccca00a078 [BugFix]Ensure that bf16 arrays are created as expected (#16436)
ccca00a078 is described below

commit ccca00a078112128aaf4e0c2714273b4a1782a2c
Author: Bin Li <43...@users.noreply.github.com>
AuthorDate: Sun Jan 21 04:07:40 2024 +0800

    [BugFix]Ensure that bf16 arrays are created as expected (#16436)
    
    Co-authored-by: Bin Li <bi...@amd.com>
---
 python/tvm/runtime/ndarray.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
index d6902bc625..300311d03b 100644
--- a/python/tvm/runtime/ndarray.py
+++ b/python/tvm/runtime/ndarray.py
@@ -176,6 +176,8 @@ class NDArray(NDArrayBase):
         if (not source_array.flags["C_CONTIGUOUS"]) or (
             dtype == "bfloat16" or dtype != np_dtype_str
         ):
+            if dtype == "bfloat16":
+                source_array = np.frombuffer(source_array.tobytes(), "uint16")
             source_array = np.ascontiguousarray(
                 source_array, dtype="uint16" if dtype == "bfloat16" else dtype
             )