From 3a3acde817ffa1cb298a346fcde963d7e25f6a4e Mon Sep 17 00:00:00 2001
From: Erika Hunhoff <erika.hunhoff@amd.com>
Date: Mon, 16 Sep 2024 15:09:51 -0600
Subject: [PATCH] try bfloat16 from tensorflow

---
 .../basic/matrix_multiplication/single_core/aie2.py          | 4 +++-
 python/extras/util.py                                        | 4 +++-
 python/requirements.txt                                      | 4 ++--
 python/requirements_ml.txt                                   | 2 --
 utils/mlir_aie_wheels/requirements.txt                       | 5 ++---
 5 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/programming_examples/basic/matrix_multiplication/single_core/aie2.py b/programming_examples/basic/matrix_multiplication/single_core/aie2.py
index e7feedaba4..a4b0e3fa89 100644
--- a/programming_examples/basic/matrix_multiplication/single_core/aie2.py
+++ b/programming_examples/basic/matrix_multiplication/single_core/aie2.py
@@ -6,10 +6,12 @@
 # (c) Copyright 2023 AMD Inc.
 
 import argparse
-from bfloat16 import bfloat16
 import numpy as np
+import tensorflow as tf
 import sys
 
+bfloat16 = tf.bfloat16.as_numpy_dtype
+
 from aie.dialects.scf import yield_, for_ as range_
 from aie.dialects.aiex import npu_dma_memcpy_nd, npu_sync
 
diff --git a/python/extras/util.py b/python/extras/util.py
index de42f42801..0a1fe328f1 100644
--- a/python/extras/util.py
+++ b/python/extras/util.py
@@ -1,5 +1,4 @@
 from collections import defaultdict
-from bfloat16 import bfloat16
 import contextlib
 import ctypes
 import inspect
@@ -12,8 +11,11 @@
 from pathlib import Path
 from typing import Callable, List, Optional, Sequence, Tuple, Union
 
+import tensorflow as tf
 import numpy as np
 
+bfloat16 = tf.bfloat16.as_numpy_dtype
+
 from .meta import op_region_builder
 from ..extras import types as T
 from ..ir import (
diff --git a/python/requirements.txt b/python/requirements.txt
index fa2ab6391a..cb304a777f 100644
--- a/python/requirements.txt
+++ b/python/requirements.txt
@@ -4,11 +4,11 @@ cmake==3.27.9
 dataclasses>=0.6, <=0.8
 filelock==3.13.1
 lit
-numpy<2
+numpy>=1.19.5, <=1.26
 pandas
 psutil
 pybind11>=2.9.0, <=2.10.3
 rich
 setuptools
-bfloat16
 wheel
+tensorflow-cpu
\ No newline at end of file
diff --git a/python/requirements_ml.txt b/python/requirements_ml.txt
index 164e0a8c53..2c5b613ae7 100644
--- a/python/requirements_ml.txt
+++ b/python/requirements_ml.txt
@@ -1,4 +1,2 @@
 --index-url https://download.pytorch.org/whl/cpu
-numpy<2
 torch
-bfloat16
diff --git a/utils/mlir_aie_wheels/requirements.txt b/utils/mlir_aie_wheels/requirements.txt
index ee4cb3d1ec..0e32e47633 100644
--- a/utils/mlir_aie_wheels/requirements.txt
+++ b/utils/mlir_aie_wheels/requirements.txt
@@ -2,10 +2,9 @@ cmake==3.27.9
 dataclasses
 importlib-metadata
 ninja
-numpy<2
+numpy>=1.19.5, <=1.26
 pip
 pybind11[global]>=2.10.4
 rich
 setuptools>=42
-wheel
-bfloat16
\ No newline at end of file
+wheel
\ No newline at end of file